BZOJ 4175 小G的电话本 ——NTT

后缀自动机统计出现了各种次数的串的和。

就是所谓的生成函数

然后FFT卷积即可。

卷积快速幂$n\log n \log n$

注意一下实现，可以少两次NTT

#include <map>
#include <cmath>
#include <queue>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
#define F(i,j,k) for (int i=j;i<=k;++i)
#define D(i,j,k) for (int i=j;i>=k;--i)
#define ll long long
#define mp make_pair
#define maxn 500005
#define md 1005060097
#define g 5
 
int A[maxn],B[maxn];
 
int rev[maxn];
const double pi=acos(-1.0);
 
int ksm(int a,int b)
{
    int ret=1;
    for (;b;b>>=1,a=(ll)a*a%md) if (b&1) ret=(ll)ret*a%md;
    return ret;
}
 
void NTT(int * x,int n,int flag)
{
    F(i,0,n-1) if (rev[i]>i) swap(x[rev[i]],x[i]);
    for (int m=2;m<=n;m<<=1)
    {
        int wn=ksm(g,((md-1)/m*flag+md-1)%(md-1));
        for (int i=0;i<n;i+=m)
        {
            int w=1;
            for (int j=0;j<(m>>1);++j)
            {
                int u=x[i+j],v=(ll)x[i+j+(m>>1)]*w%md;
                x[i+j]=(u+v)%md; x[i+j+(m>>1)]=(u-v+md)%md;
                w=(ll)w*wn%md;
            }
        }
    }
    if (flag==-1)
    {
        int inv=ksm(n,md-2);
        F(i,0,n-1) x[i]=(ll)x[i]*inv%md;
    }
}
 
int k,m,n,N,len;
 
struct Suffix_Automaton{
    int cnt,last,l[maxn],fa[maxn],go[maxn][26],siz[maxn];
    int v[maxn],q[maxn];
    char s[maxn];
    void init(){cnt=last=1;memset(go,0,sizeof go);}
    void add(int x)
    {
        int p=last,np=last=++cnt; l[np]=l[p]+1;
        for (;p&&!go[p][x];p=fa[p]) go[p][x]=np;
        if (!p) fa[np]=1;
        else
        {
            int q=go[p][x];
            if (l[q]==l[p]+1) fa[np]=q;
            else
            {
                int nq=++cnt; l[nq]=l[p]+1;
                memcpy(go[nq],go[q],sizeof go[q]);
                fa[nq]=fa[q];
                fa[q]=fa[np]=nq;
                for (;p&&go[p][x]==q;p=fa[p]) go[p][x]=nq;
            }
        }
        siz[last]++;
    }
    void radix()
    {
        F(i,1,cnt) v[l[i]]++;
        F(i,1,cnt) v[i]+=v[i-1];
        F(i,1,cnt) q[v[l[i]]--]=i;
        D(i,cnt,1) siz[fa[q[i]]]+=siz[q[i]];
    }
    void solve()
    {
        scanf("%s",s); int n=strlen(s);
        init(); F(i,0,n-1) add(s[i]-'a');
        radix();
    }
}SAM;
 
int main()
{
    scanf("%d%d",&k,&m); SAM.solve();
    F(i,2,SAM.cnt) A[SAM.siz[i]]=(A[SAM.siz[i]]+(ll)(SAM.l[i]-SAM.l[SAM.fa[i]])*SAM.siz[i]%md)%md;B[0]++;
    for (n=1,len=0;n<=2*m;n<<=1) len++;
    F(i,0,n-1) rev[i]=(rev[i>>1]>>1)|((i&1)<<(len-1));
    F(i,m+1,n-1) B[i]=A[i]=0;
    for (;k;k>>=1)
    {
        NTT(A,n,1);
        if (k&1)
        {
            NTT(B,n,1);
            F(i,0,n-1) B[i]=(ll)B[i]*A[i]%md;
            NTT(B,n,-1);F(i,m+1,n-1) B[i]=0;
        }
        F(i,0,n-1) A[i]=(ll)A[i]*A[i]%md;
        NTT(A,n,-1);F(i,m+1,n-1)A[i]=0;
    }
    printf("%d\n",(B[m]%md+md)%md);
}