NTT

DP_PTSD發表於2024-10-27

NTT

線性卷積

定義:

\[(f * g)[i] = \sum_{j=0}^{i} f[j] \cdot g[i-j] \]

卷積定理:

\[\mathcal{F}(f * g) = \mathcal{F}(f) \cdot \mathcal{F}(g) \]

於是,求線性卷積可以轉化為,先變換,再直接相乘,最後逆變換。

注意:序列長度須變成N+M-1,對齊2^k。

迴圈卷積(圓周卷積)

一般用於兩個等長序列。

定義:

\[(f * g)[i] = \sum_{j=0}^{N-1} f[j] \cdot g[i-j] \]

迴圈卷積可以透過線性卷積來求解。設 \(B[i]\) 為線性卷積結果,
則迴圈卷積 \(C[i]=B[i]+B[i+N]\)

大模數NTT

適用於答案為整數,範圍4e18以內,有效彌補FFT在大於1e15後精度不足的缺陷。

__int128常數較大,若答案範圍較小(9e8以內),則用小模數或FFT

struct Polynomial
{
    static const ll mod = 4179340454199820289LL; // 998244353LL
    vector<ll> z;
    vector<int> r;
    Polynomial(vector<ll> &a):z(a)
    {
        int n = a.size();
        r.resize(n);
        for (int i=0; i<n; i++)
            r[i] = (i&1)*(n/2) + r[i/2]/2;
        ntt(z, n, 1);
    }
    ll power(ll a, ll b)
    {
        ll res = 1;
        for(; b; b>>=1, a = (__int128)a * a % mod)
            if(b&1) res = (__int128)res * a % mod;
        return res;
    }
    void ntt(vector<ll> &a, int n, int opt)
    {
        for(int i=0; i<n; i++)
            if (r[i]<i) swap(a[i], a[r[i]]);
        
        for(int k=2; k<=n; k*=2)
        {
            ll gn = power(3, (mod-1)/k);
            for(int i=0; i<n; i+=k)
            {
                ll g = 1;
                for(int j=0; j<k/2; j++, g = (__int128)g * gn % mod)
                {
                    ll t = (__int128)a[i+j+k/2] * g % mod;
                    a[i+j+k/2] = (a[i+j] + mod - t) % mod;
                    a[i+j] = (a[i+j] + t) % mod;
                }
            }
        }
        if (opt == -1)
        {
            reverse(a.begin()+1, a.end());
            ll inv = power(n, mod-2);
            for(int i=0; i<n; i++) a[i] = (__int128)a[i] * inv % mod;    
        }
    }
};