快速数论变换 NTT

更新日志 2025/01/08：开工。

2025/01/09：更新求逆元代码，给出整合封装模板。

前置知识

本文将在 \(\text{FFT}\) 基础上进行讲解。

同时请确保会 欧拉函数 等数论基础。

作用

\(\text{FFT}\) 需要用到复数，而 double 使我们面临严重的精度问题。

所以，我们需要一个更精确的算法——\(\text{NTT}\)。

这个算法的核心，就是用原根代替单位根。

原根

阶

若 \(a,p\) 互素且 \(p>1\)，使 \(a^n\equiv1\pmod p\) 的最小正整数 \(n\)，称之为 \(a\) 模 \(p\) 的阶，记作 \(\delta_p(a)\)。

如 \(\delta_7(2)=3\)，因为 \(2^3\equiv1\pmod7\)。

原根

若 \(p\in Z^*,a\in Z\)，且 \(\delta_p(a)=\varphi(p)\)，则称 \(a\) 为模 \(p\) 的一个原根。

原根是不唯一的。如果 \(p\) 有原根，那么它的原根个数为 \(\varphi(\varphi(p))\)。

当然这不重要，重要的是原根具有和单位根相似的性质！

我们可以得到这样一个结论：\(\omega_n^1\equiv g^\frac{p-1}{n}\pmod p\)。

然后我们只需要把 \(\text{FFT}\) 中的 \(\omega_n\) 替换掉就好了。

常见实现

一组很常见的 \(g,p\) 是 \(998244353,3\)。

逆元 \(g^{-1}=332748118\)

至于如何求任意质数的原根，我们分解 \(n-1\) 质因子，若 \(g^\frac{p-1}{p_i}\not\equiv1\pmod p\) 恒成立，\(g\) 即为 \(p\) 的原根。

多项式乘法模板

const int p=998244353,g=3,gi=332748118;
typedef vec<ll> poly;

ll qpow(ll a,ll b){
    ll res=1;
    while(b){
        if(b&1)(res*=a)%=p;
        (a*=a)%=p;
        b>>=1;
    }
    return res;
}

vec<int> r;
void ntt(int lim,poly &a,int type){
    a.resize(lim);
    repl(i,0,lim)if(i<r[i])swap(a[i],a[r[i]]);
    for(int mid=1;mid<lim;mid<<=1){
        ll w1=qpow((~type)?g:gi,(p-1)/(mid<<1));
        for(int j=0;j<lim;j+=(mid<<1)){
            ll w=1;
            repl(k,0,mid){
                ll x=a[j+k],y=w*a[j+mid+k]%p;
                a[j+k]=(x+y)%p;
                a[j+mid+k]=(x-y+p)%p;
                w=w*w1%p;
            }
        }
    }
}
poly operator*(poly a,poly b){
    int lim=1,l=0,len=a.size()+b.size()-1;
    while(lim<a.size()+b.size())lim<<=1,l++;
    r.resize(lim);
    repl(i,0,lim)r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
    ntt(lim,a,1);ntt(lim,b,1);
    poly c(lim);
    repl(i,0,lim)c[i]=(a[i]*b[i])%p;
    ntt(lim,c,-1);
    c.resize(len);
    ll inv=qpow(lim,p-2);
    repl(i,0,c.size())(c[i]*=inv)%=p;
    return c;
}

例题代码

#include<bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef __int128 i128;
typedef double db;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
typedef pair<int,ll> pil;
typedef pair<ll,int> pli;
template <typename Type>
using vec=vector<Type>;
template <typename Type>
using grheap=priority_queue<Type>;
template <typename Type>
using lrheap=priority_queue<Type,vector<Type>,greater<Type> >;
#define fir first
#define sec second
#define pub push_back
#define pob pop_back
#define puf push_front
#define pof pop_front
#define chmax(a,b) a=max(a,b)
#define chmin(a,b) a=min(a,b)
#define rep(i,x,y) for(int i=(x);i<=(y);i++)
#define repl(i,x,y) for(int i=(x);i<(y);i++)
#define per(i,x,y) for(int i=(x);i>=(y);i--)

const int inf=0x3f3f3f3f;
const ll INF=0x3f3f3f3f3f3f3f3f;

int n,m;

const int p=998244353,g=3,gi=332748118;
typedef vec<ll> poly;

ll qpow(ll a,ll b){
    ll res=1;
    while(b){
        if(b&1)(res*=a)%=p;
        (a*=a)%=p;
        b>>=1;
    }
    return res;
}

vec<int> r;
void ntt(int lim,poly &a,int type){
    a.resize(lim);
    repl(i,0,lim)if(i<r[i])swap(a[i],a[r[i]]);
    for(int mid=1;mid<lim;mid<<=1){
        ll w1=qpow((~type)?g:gi,(p-1)/(mid<<1));
        for(int j=0;j<lim;j+=(mid<<1)){
            ll w=1;
            repl(k,0,mid){
                ll x=a[j+k],y=w*a[j+mid+k]%p;
                a[j+k]=(x+y)%p;
                a[j+mid+k]=(x-y+p)%p;
                w=w*w1%p;
            }
        }
    }
}
poly operator*(poly a,poly b){
    int lim=1,l=0,len=a.size()+b.size()-1;
    while(lim<a.size()+b.size())lim<<=1,l++;
    r.resize(lim);
    repl(i,0,lim)r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
    ntt(lim,a,1);ntt(lim,b,1);
    poly c(lim);
    repl(i,0,lim)c[i]=(a[i]*b[i])%p;
    ntt(lim,c,-1);
    c.resize(len);
    ll inv=qpow(lim,p-2);
    repl(i,0,c.size())(c[i]*=inv)%=p;
    return c;
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    cin>>n>>m;
    poly a(n+1),b(m+1);
    rep(i,0,n)cin>>a[i],a[i]=(a[i]%p+p)%p;
    rep(j,0,m)cin>>b[j],b[j]=(b[j]%p+p)%p;
    poly c=a*b;
    repl(i,0,c.size())cout<<c[i]<<" ";
    return 0;
}

求逆元

定义

我们对一个多项式的逆元的定义是：

对于一个多项式 \(A(x)\)，其逆元 \(B(x)\) 满足 \(A(x)\times B(x)\equiv1\pmod {x^n}\)

是什么意思呢？就是说，删去次数 \(\ge n\) 的项，除了常数项为 \(1\) 以外，其他项均被消掉。

公式

又到了喜闻乐见的推式子环节。

\[A(x)\times B(x)\equiv1\pmod{x^n} \]

设 \(B'(x)\) 满足 \(A(x)\times B'(x)\equiv1\pmod{x^{\lceil\frac{n}{2}\rceil}}\)

然后很显然的一个式子：

\[A(x)\times B(x)\equiv1\pmod{x^{\lceil\frac{n}{2}\rceil}} \]

两项相减得到：

\[B(x)-B'(x)\equiv0\pmod{x^{\lceil\frac{n}{2}\rceil}}\\ \Rightarrow B^2(x)+B'^2(x)-2B(x)B'(x)\equiv0\pmod{x^n} \]

在这里我们停一停，你会发现，模数又变回了 \(x^n\)，为什么呢？

\[\because F(x)=\sum_{i=0}^na_ix^i\equiv0\pmod{x^{\lceil\frac{n}{2}\rceil}}\\ \therefore \forall i\in[0,n],a_i=0 \]

接下来我们考虑平方后，\(x\) 的前 \(n\) 项：

\[a_i'=\sum_{j=0}^ia_ja_{i-j} \]

这两项中至少有一项属于 \([0,{\lceil\frac{n}{2}\rceil}]\)，故而系数为 \(0\)。

言归正传，我们在上式两侧同时乘 \(A(x)\) 得到：

\[B(x)+B'^2(x)A(x)-2B'(x)\equiv0\pmod{x^n} \]

注意在模 \(x^n\) 情况下 \(A(x)B'(x)\) 不同余于 \(1\)。

我们最后移一次项：

\[B(x)\equiv2B'(x)-A(x)B'^2(x)\pmod{x^n} \]

你发现，我们可以使用递归的方式求解 \(B(x)\)。具体的，我们已知 \(A(x)\)，只需要得到 \(B'(x)\)，就可以得到 \(B(x)\) 了。

实现

通常情况下，得到的系数都会对一个特殊模数 \(p\) 取余。因此，我们可以使用 \(\text{NTT}\) 求解，\(\text{NTT}\) 的模数也设为 \(p\) 即可。

当只剩下一个常数项时，\(B\) 内唯一的项显然就是那个常数的逆元。

否则，我们递归得到 \(B'(x)\)。

接下来，我们使用 \(\text{NTT}\)，获取 \(A(x)\) 和 \(B'(x)\)，并利用上方公式求得 \(B(x)\) 即可。

注意事项

进行操作的多项式应只保留需要部分，不然会导致超时！

模板

poly getinv(poly a,int len){
    if(len==1)return poly(1,qpow(a[0],p-2));
    poly aa=a;aa.resize(len);
    poly bb=getinv(a,len+1>>1);
    poly b=poly(1,2)*bb-aa*bb*bb;
    b.resize(len);
    return b;
}

例题代码

#include<bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef unsigned long long ull;
typedef __int128 i128;
typedef double db;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
typedef pair<int,ll> pil;
typedef pair<ll,int> pli;
template <typename Type>
using vec=vector<Type>;
template <typename Type>
using grheap=priority_queue<Type>;
template <typename Type>
using lrheap=priority_queue<Type,vector<Type>,greater<Type> >;
#define fir first
#define sec second
#define pub push_back
#define pob pop_back
#define puf push_front
#define pof pop_front
#define chmax(a,b) a=max(a,b)
#define chmin(a,b) a=min(a,b)
#define rep(i,x,y) for(int i=(x);i<=(y);i++)
#define repl(i,x,y) for(int i=(x);i<(y);i++)
#define per(i,x,y) for(int i=(x);i>=(y);i--)

const int inf=0x3f3f3f3f;
const ll INF=0x3f3f3f3f3f3f3f3f;
const int mod=998244353;

namespace NTT{
    const int p=998244353,g=3,gi=332748118;
    typedef vec<ll> poly;
    ll qpow(ll a,ll b){
        ll res=1;
        while(b){
            if(b&1)(res*=a)%=p;
            (a*=a)%=p;
            b>>=1;
        }
        return res;
    }

    vec<int> r;
    void ntt(int lim,poly &a,int type){
        a.resize(lim);
        repl(i,0,lim)if(i<r[i])swap(a[i],a[r[i]]);
        for(int mid=1;mid<lim;mid<<=1){
            ll w1=qpow((~type)?g:gi,(p-1)/(mid<<1));
            for(int j=0;j<lim;j+=(mid<<1)){
                ll w=1;
                repl(k,0,mid){
                    ll x=a[j+k],y=w*a[j+mid+k]%p;
                    a[j+k]=(x+y)%p;
                    a[j+mid+k]=(x-y+p)%p;
                    w=w*w1%p;
                }
            }
        }
    }
    poly operator*(poly a,poly b){
        int lim=1,l=0,len=a.size()+b.size()-1;
        while(lim<a.size()+b.size())lim<<=1,l++;
        r.resize(lim);
        repl(i,0,lim)r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
        ntt(lim,a,1);ntt(lim,b,1);
        repl(i,0,lim)a[i]=(a[i]*b[i])%p;
        ntt(lim,a,-1);
        a.resize(len);
        ll inv=qpow(lim,p-2);
        repl(i,0,a.size())(a[i]*=inv)%=p;
        return a;
    }
    poly operator+(poly a,poly b){
        if(a.size()<b.size())a.resize(b.size());
        repl(i,0,b.size())a[i]=(a[i]+b[i])%p;
        return a;
    }
    poly operator-(poly a,poly b){
        if(a.size()<b.size())a.resize(b.size());
        repl(i,0,b.size())a[i]=(a[i]-b[i]+p)%p;
        return a;
    }
}
using namespace NTT;

poly getinv(poly a,int len){
    if(len==1)return poly(1,qpow(a[0],p-2));
    poly aa=a;aa.resize(len);
    poly bb=getinv(a,len+1>>1);
    poly b=poly(1,2)*bb-aa*bb*bb;
    b.resize(len);
    return b;
}

int n;
poly a;

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);
    cin>>n;a.resize(n);
    repl(i,0,n)cin>>a[i],a[i]%=p;
    poly b=getinv(a,n);
    repl(i,0,b.size())cout<<b[i]<<" ";
    return 0;
}

整合封装模板

namespace NTT{
    const int p=998244353,g=3,gi=332748118;
    typedef vec<ll> poly;
    ll qpow(ll a,ll b){
        ll res=1;
        while(b){
            if(b&1)(res*=a)%=p;
            (a*=a)%=p;
            b>>=1;
        }
        return res;
    }

    vec<int> r;
    void ntt(int lim,poly &a,int type){
        a.resize(lim);
        repl(i,0,lim)if(i<r[i])swap(a[i],a[r[i]]);
        for(int mid=1;mid<lim;mid<<=1){
            ll w1=qpow((~type)?g:gi,(p-1)/(mid<<1));
            for(int j=0;j<lim;j+=(mid<<1)){
                ll w=1;
                repl(k,0,mid){
                    ll x=a[j+k],y=w*a[j+mid+k]%p;
                    a[j+k]=(x+y)%p;
                    a[j+mid+k]=(x-y+p)%p;
                    w=w*w1%p;
                }
            }
        }
    }
    poly operator*(poly a,poly b){
        int lim=1,l=0,len=a.size()+b.size()-1;
        while(lim<a.size()+b.size())lim<<=1,l++;
        r.resize(lim);
        repl(i,0,lim)r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
        ntt(lim,a,1);ntt(lim,b,1);
        repl(i,0,lim)a[i]=(a[i]*b[i])%p;
        ntt(lim,a,-1);
        a.resize(len);
        ll inv=qpow(lim,p-2);
        repl(i,0,a.size())(a[i]*=inv)%=p;
        return a;
    }
    poly operator+(poly a,poly b){
        if(a.size()<b.size())a.resize(b.size());
        repl(i,0,b.size())a[i]=(a[i]+b[i])%p;
        return a;
    }
    poly operator-(poly a,poly b){
        if(a.size()<b.size())a.resize(b.size());
        repl(i,0,b.size())a[i]=(a[i]-b[i]+p)%p;
        return a;
    }
}
using namespace NTT;
posted @   LastKismet  阅读(70)  评论(1)  编辑  收藏  举报
相关博文：
·  快速傅里叶变换 FFT
·  分治FFT
·  【学习笔记】- 基础数论：快速数论变换
·  快速数论变换 | NTT 初学
·  NTT（快速数论变换）
阅读排行：
· 地球OL攻略 —— 某应届生求职总结
· 周边上新：园子的第一款马克杯温暖上架
· Open-Sora 2.0 重磅开源！
· 提示词工程——AI应用必不可少的技术
· .NET周刊【3月第1期 2025-03-02】
点击右上角即可分享
微信分享提示
AI IDE Trae