← all categories

Ntt a6e15b2d

Number Theoretic Transform (in-place, iterative) and exact polynomial multiplication mod `mod` (must be a prime of the form c*2^k+1 with a known primitive root). Default mod = (483<<21)+1, root = 62; other options: (119<<21)+1 (=998244353, root 3, most common), 5<<25+1, 7<<26+1. Source: kactl (CC0), see NOTICE.md.

Time: O(n \log n). kactl-derived, tested

content/math/ntt.h

const ll ntt_mod = (483LL << 21) + 1, ntt_root = 62;

ll ntt_modpow(ll a, ll b, ll mod) {
    ll r = 1;
    a %= mod;
    while (b) {
        if (b & 1) r = r * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return r;
}

void ntt(vector<ll> &a) {
    int n = (int)a.size(), L = 31 - __builtin_clz(n);
    static vector<ll> rt(2, 1);
    for (static int k = 2, s = 2; k < n; k *= 2, s++) {
        rt.resize(n);
        ll z[] = {1, ntt_modpow(ntt_root, ntt_mod >> s, ntt_mod)};
        for (int i = k; i < 2 * k; i++) rt[i] = rt[i / 2] * z[i & 1] % ntt_mod;
    }
    vector<int> rev(n);
    for (int i = 0; i < n; i++) rev[i] = (rev[i / 2] | (i & 1) << L) / 2;
    for (int i = 0; i < n; i++) if (i < rev[i]) swap(a[i], a[rev[i]]);
    for (int k = 1; k < n; k *= 2)
        for (int i = 0; i < n; i += 2 * k)
            for (int j = 0; j < k; j++) {
                ll z = rt[j + k] * a[i + j + k] % ntt_mod, &ai = a[i + j];
                a[i + j + k] = ai - z + (z > ai ? ntt_mod : 0);
                ai += (ai + z >= ntt_mod ? z - ntt_mod : z);
            }
}

vector<ll> conv(const vector<ll> &a, const vector<ll> &b) {
    if (a.empty() || b.empty()) return {};
    int s = (int)a.size() + (int)b.size() - 1, B = 32 - __builtin_clz(s), n = 1 << B;
    ll inv = ntt_modpow(n, ntt_mod - 2, ntt_mod);
    vector<ll> L(a), R(b), out(n);
    L.resize(n); R.resize(n);
    ntt(L); ntt(R);
    for (int i = 0; i < n; i++) out[-i & (n - 1)] = L[i] * R[i] % ntt_mod * inv % ntt_mod;
    ntt(out);
    return {out.begin(), out.begin() + s};
}