洛谷 P4199 万径人踪灭 回文树+fft

题目背景

题目描述


输入输出格式

输入格式：
一行，一个只包含a,b两种字符的字符串

输出格式：
一行，一个整数表示问题的答案

输入输出样例

输入样例#1：
abaabaa
输出样例#1：
14
输入样例#2：
aaabbbaaa
输出样例#2：
44
输入样例#3：
aaaaaaaa
输出样例#3：
53
说明

分析：
可以考虑把总的对称的方案数减去连成一块的方案数。
显然后面一个就是回文串个数，直接回文树解决。
假设我们的对称中心是$k$，那么每个满足的方案数$2^{x}-1$，其中$x$为左右两边配对的对数，每对可以选或不选，减一表示不能为空。
然后我发现这个和gdkoi那道基站是一个道理，不知道就跳过……
那么$k$的答案为

$$\sum_{i+j=2k}[s[i]==s[j]]$$
因为只有两种字符，考虑每一种字符，设 $f[k]$为 $a$对数对$k$的方案，
则
$$f[k]=\sum_{i+j=2*k}[s[i]==a]\ and\ [s[j]==a]$$
假设 $x[i]=(a[i]==a)$
那么
$$f[k]=\sum_{i+j=2*k}x[i]*x[j]$$
这个就是一个卷积形式，直接fft。
发现每个对会算两次，而在中心的字符只会算一次，判断解决即可。

代码：

// luogu-judger-enable-o2
#include <iostream>
#include <cstdio>
#include <cstring>
#include <cmath>
#define LL long long

const int maxn=3e5+7;
const double pi=acos(-1);
const LL mod=1e9+7;

using namespace std;

LL n,cnt,len;
LL r[maxn];
char s[maxn];
LL f[maxn],g[maxn],ans;

struct node{
    LL fail,len,sum;
    LL son[2];
}t[maxn];

struct rec{
    double x,y;
}w[maxn],a[maxn],b[maxn],c[maxn];

rec operator +(rec a,rec b)
{
    return (rec){a.x+b.x,a.y+b.y};
}

rec operator -(rec a,rec b)
{
    return (rec){a.x-b.x,a.y-b.y};
}

rec operator *(rec a,rec b)
{
    return (rec){a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x};
}

rec operator !(rec a)
{
    return (rec){a.x,-a.y};
}

void build_tree()
{
    cnt=1;
    t[0].fail=1;
    t[0].len=0;
    t[1].fail=0;
    t[1].len=-1;
    LL now=1;
    for (LL i=0;i<n;i++)
    {
        while (s[i]!=s[i-t[now].len-1]) now=t[now].fail;
        if (!t[now].son[s[i]-'a'])
        {
            cnt++;
            LL k=t[now].fail;
            while (s[i]!=s[i-t[k].len-1]) k=t[k].fail;
            t[cnt].fail=t[k].son[s[i]-'a'];
            t[now].son[s[i]-'a']=cnt;
            t[cnt].len=t[now].len+2;
            t[cnt].sum=(t[t[cnt].fail].sum+1)%mod;
        }
        now=t[now].son[s[i]-'a'];
        ans=(ans+t[now].sum)%mod;
    }
}

void fft(rec *a,LL f)
{
    for (LL i=0;i<len;i++)
    {
        if (i<r[i]) swap(a[i],a[r[i]]);
    }
    w[0]=(rec){1,0};
    for (LL i=2;i<=len;i*=2)
    {
        rec wn=(rec){cos(2*pi/i),f*sin(2*pi/i)};
        for (LL j=i/2;j>=0;j-=2) w[j]=w[j/2];
        for (LL j=1;j<i/2;j+=2) w[j]=w[j-1]*wn;
        for (LL j=0;j<len;j+=i)
        {
            for (LL k=0;k<i/2;k++)
            {
                rec u=a[j+k],v=a[j+k+i/2]*w[k];
                a[j+k]=u+v;
                a[j+k+i/2]=u-v;
            }
        }
    }
}

void init(LL len)
{
    LL k=trunc(log(len+0.5)/log(2));
    for (LL i=0;i<len;i++)
    {
        r[i]=(r[i>>1]>>1)|((i&1)<<(k-1));
    }
}

void FFT(LL *x,LL *y,LL *z,LL n,LL m)
{   
    len=1;
    while (len<(n+m-1)) len*=2;
    init(len);
    for (LL i=0;i<len;i++)
    {
        LL A,B;
        if (i<n) A=x[i]; else A=0;
        if (i<m) B=y[i]; else B=0;
        a[i]=(rec){A,0};
        b[i]=(rec){B,0};
    }
    fft(a,1); fft(b,1);
    for (int i=0;i<len;i++) c[i]=a[i]*b[i];
    fft(c,-1);
    for (int i=0;i<len;i++) z[i]=(LL)(c[i].x/len+0.5);
}

LL power(LL x,LL y)
{
    if (y==0) return 1;
    if (y==1) return x;
    LL c=power(x,y/2);
    c=(c*c)%mod;
    if (y%2) c=(c*x)%mod;
    return c;
}

int main()
{
    scanf("%s",s);
    n=strlen(s);
    for (LL i=0;i<n;i++)
    {
        f[i]=(s[i]=='a');
        g[i]=(s[i]=='b');
    }       
    build_tree();   
    FFT(f,f,f,n,n); 
    FFT(g,g,g,n,n); 
    ans=mod-ans;    
    for (LL i=0;i<=2*(n-1);i++)
    {
        if (i%2==0)
        {
            f[i]+=(s[i/2]=='a');
            g[i]+=(s[i/2]=='b');
        }
        f[i]/=2;
        g[i]/=2;
        ans=(ans+power(2,f[i]+g[i])+mod-1)%mod;
    }
    printf("%lld",ans);
}