快速傅里叶变换（FFT）

GYM 101002.E K-Inversions

#include<bits/stdc++.h>
using namespace std;
const double PI=acos(-1.0);
typedef complex<double> cp;
const int maxn=3e6+100;
cp arr[maxn],b[maxn];
cp omega[maxn],inv[maxn];
int eg[maxn];
int n=1;
void init()
{
    for(int i=0;i<n;++i)
    {
        omega[i]=cp(cos(2*PI*i/n),sin(2*PI*i/n));
        inv[i]=conj(omega[i]);
    }
}
void FFT(cp *a,cp *omg)
{
    for(int i=0;i<n;++i)
    {
        if(i<eg[i])
            swap(a[i],a[eg[i]]);
    }
    for(int l=2;l<=n;l<<=1)
    {
        int m=l/2;
        for(cp *p=a;p!=a+n;p+=l)
        {
            for(int i=0;i<m;++i)
            {
                cp t=omg[n/l*i]*p[i+m];
                p[i+m]=p[i]-t;
                p[i]+=t;
            }
        }
    }
}
int main()
{
    string str;
    cin>>str;
    int len=str.size();
    while(n<=len)
        n<<=1;
        n<<=1;
        init();
    for(int i=0;i<len;++i)
    {
        if(str[i]=='A')
            arr[i].real(1);
        else
            b[len-i-1].real(1);
    }
    int limit=0;
    while((1<<limit)<n)
        ++limit;
    for(int i=0;i<n;++i)
    {
        int t=0;
        for(int j=0;j<limit;++j)
        {
            if((i>>j)&1)
                t|=(1<<(limit-1-j));
        }
        eg[i]=t;
    }
    FFT(arr,omega);
    FFT(b,omega);
    for(int i=0;i<n;++i)
    {
        arr[i]*=b[i];
    }
    FFT(arr,inv);
    for(int i=len;i<len*2-1;++i)
    {
        int ans=(int)(arr[i].real()/n+0.5);
        cout<<ans<<endl;
    }
    return 0;
}