ABC281 - Ex を非再帰で書いた

タイトルのままです 再帰アレルギーの方のために

じゃなくて普通にこっち先に思い付いた

方針は公式解説と同じです 別にそんな速くはなさそう

二冪の prefix にFFT のマージテクでできたやつかけていけばいいよねってだけ

自分の NTT でやったら 661ms で ACL だと 457 ms で悲しくなりました 書き直したいです

コード

#include <bits/stdc++.h>
#include <atcoder/convolution>
#include <atcoder/modint>

using namespace std;
using namespace atcoder;

using mint = modint998244353;
constexpr int mod = 998244353;

int p2(int n){
  while(n & (n - 1)) ++n;
  return n;
}

vector<mint> pre(const vector<mint> &a, int n){
  return {begin(a), begin(a) + n};
}

int main() {
  int n, a; cin >> n >> a;

  if(n == 1){
    cout << a % mod << "\n";
    return 0;
  }

  vector<mint> c(n + 1, 1), inv(n + 1, 1);

  for(int i = 2; i <= n; ++i) inv[i] = -inv[mod % i] * (mod / i);

  for(int i = 0; i < n; ++i) c[i + 1] = c[i] * (a - i) * inv[i + 1];

  c.erase(begin(c), begin(c) + 2);
  c.resize(p2(c.size()));

  stack<vector<mint>> sf, sa; sf.push(c);

  for(int i = 0; i < n - 1; ++i){
    vector<mint> f = sf.top();
    while(int(f.size()) > 1){
      sf.push(f = pre(f, f.size() / 2));
    }

    mint x = sf.top()[0]; sf.pop();

    if(i == n - 2){
      cout << x.val() << "\n";
      return 0;
    }
    
    f = {1, x};
    while(!sa.empty()){
      vector<mint> g = sa.top();
      if(f.size() != g.size()) break;
      sa.pop();
      f = convolution(f, g);
    }
    sa.push(f);

    int m = int(sf.top().size());
    f = convolution(f, sf.top()); sf.pop();
    f = pre(f, m);
    sf.emplace(begin(f) + (m/2), end(f));
  }
}

提出結果

atcoder.jp