数論変換

整数環FFTとも

基本的にはFFTと同じだが、あちらが実数を使って計算するのに対しこちらは有理数を素数で割った余りを使って計算する

誤差がないというメリットがあるが、\(MOD\)素数より大きい整数を扱えないというデメリットがある

ここで用いる素数は、\(P-1\)が素因数に\(2\)を多く含むような\(P\)であることが望ましい(例えばNTTでよく使われる\(998244353\)は\(119 \cdot 2^{23}+1\)である)

計算量は\(N=max(\vert A \vert,\vert B \vert)\)として、\(Ο(NlogNlogMOD)\)

github

class number_theoretic_transform{
    // Copyright (c) 2023 0214sh7
    // https://github.com/0214sh7/library/
    private:
    
    const long long MOD = 998244353;
    const long long k = 119;
    const long long n = 23;
    const long long pri_root = 3;
    //MOD=k*2^n+1
    
    long long BE(long long b,long long e){
        long long r=1;
        while(e){
            if(e&1){
                r=(r*b)%MOD;
            }
            b=(b*b)%MOD;
            e >>=1;
        }
        return r;
    }
    
    std::vector<long long> DFT(std::vector<long long> A){
        int N=A.size();
        if(N==1)return A;
        std::vector<long long> A0(N/2),A1(N/2),iA0,iA1,iA(N);
        for(int i=0;i<N;i++){
            if(i%2==0){
                A0[i/2]=A[i];
            }else{
                A1[i/2]=A[i];
            }
        }
        iA0=DFT(A0);
        iA1=DFT(A1);
        
        long long omega=BE(pri_root,k*(1<<n)/N);
        long long ith_zeta = 1;
        for(int i=0;i<N;i++){
            iA[i]=(iA0[i%(N/2)]+ ith_zeta*iA1[i%(N/2)])%MOD;
            ith_zeta = (ith_zeta*omega)%MOD;
        }
        return iA;
    }
     
    std::vector<long long> iDFT(std::vector<long long> iA){
        int N=iA.size();
        long long N_inverse = BE(N,MOD-2);
        std::vector<long long> A,dA,rA;
        dA=DFT(iA);
        for(int i=0;i<N;i++){
            rA.push_back(dA[(N-i)%N]);
            A.push_back((rA[i]*N_inverse)%MOD);
        }
        return A;
    }
    
    
    public:
    
    std::vector<long long> convolution(std::vector<long long> A,std::vector<long long> B){
        int deg = A.size() + B.size() -1;
        long long N=1;
        while(N<deg){N<<=1;}
        A.resize(N);B.resize(N);
        for(int i=0;i<A.size();i++){
            A[i]%=MOD;
        }
        for(int i=0;i<B.size();i++){
            B[i]%=MOD;
        }
        std::vector<long long> C(N),iC(N),iA,iB;
        iA=DFT(A);iB=DFT(B);
        for(int i=0;i<N;i++){
            iC[i]=(iA[i]*iB[i])%MOD;
        }
        C=iDFT(iC);
        C.resize(deg);
        return C;
    }
    
};

使用例

---

実行コード

#include <bits/stdc++.h>

class number_theoretic_transform{/*省略*/};

int main(void){
    
    number_theoretic_transform NTT;
    std::vector<long long> A={1,2,3},B={1,3,5},C;
    
    C = NTT.convolution(A,B);
    for(int i=0;i<C.size();i++){
        std::cout << C[i] << " ";
    }
    std::cout << std::endl;
    
    return 0;
}

出力

1 5 14 19 15 

更新履歴

---

日時	内容
2023/06/29	ライセンスのコメントアウトを変更
2021/03/25	使用例を追加/vectorをstd::vectorに変更
2020/04/04	数論変換を追加