Codechef Coders’Legacy 2018 CLSUMG Sum of Primes

Description

• 设 $f(x)$ 表示把 $x$ 拆分成两个质数之和的方案数。

• 例如 $f(10)=2$ , 因为 $10=3+7=5+5$

• $T$ 次询问，每次给出 $n$，问有多少对 $(a,b)$ 满足 $0≤a≤b＜n$ 且 $f(a)+f(b)=f(n)$ 。

• 数据范围：$1\le T≤5，0≤n≤1,000,000$ 。

Sample Input:

3
1
5
10

Sample Output:

1
4
21

Solution

• 我们把质数看做 1 ，非质数看成 0 。

• 即设一个数组 $h[i]$ ，若 $i$ 为质数则 $h[i]=1$ ，否则 $h[i]=0$ 。

• 而 $n\le10^6$ ，我们可以直接 $O(n)$ 筛出 $h[i]$ 。

• 观察 $f[i]$ 的定义和数论卷积的定义，我们可以愉快地发现 $f$ 就是两个 $h$ 卷起来！

• 之后再去一下重就能预处理出 $f$ 了！用 FFT 快速算 $f[i]$ 可以做到 $O(n\ log\ n)$ 。

• 接下来就可以轻松地处理询问了。

• 用一个桶直接 $O(n)$ 扫一遍就可以啦！

• 倒着枚举 $i$ ，则答案加上 $t[f[n]-f[i]]$ ，之后 $t[f[i]]++$ 。

• 总时间复杂度 $O(n\ log\ n)$ 。

Code

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;
typedef long long LL;
const int N=1e6+5;
const double Pi=acos(-1.0);
struct comp
{
    double r,i;
    comp(){}
    comp(double rr,double ii){r=rr,i=ii;}
    comp operator +(const comp &x)const
    {
        return comp(r+x.r,i+x.i);
    }
    comp operator -(const comp &x)const
    {
        return comp(r-x.r,i-x.i);
    }
    comp operator *(const comp &x)const
    {
        return comp(r*x.r-i*x.i,i*x.r+r*x.i);
    }
}f1[N<<2],f2[N<<2];
int T,n,m;
int f[N],g[N],h[N],rev[N<<2],tt[N];
bool bz[N];
inline void FFT(comp *y,int ff)
{
    for(int i=0;i<m;i++)
        if(i<rev[i]) swap(y[i],y[rev[i]]);
    for(int h=2;h<=m;h<<=1)
    {
        comp wn(cos(2*Pi/h),ff*sin(2*Pi/h));
        for(int i=0;i<m;i+=h)
        {
            comp w(1,0);
            for(int k=i;k<i+h/2;k++)
            {
                comp u=y[k],t=w*y[k+h/2];
                y[k]=u+t;
                y[k+h/2]=u-t;
                w=w*wn;
            }
        }
    }
    if(ff==-1) for(int i=0;i<m;i++) y[i].r/=m;
}
int main()
{
    for(int i=2;i<N;i++)
    {
        if(!bz[i]) g[++g[0]]=i;
        for(int j=1;j<=g[0] && i*g[j]<N;j++)
        {
            bz[i*g[j]]=true;
            if(i%g[j]==0) break;
        }
    }
    for(int i=2;i<N;i++) h[i]=!bz[i];
    int l=0;
    for(m=1;m<N+N;m<<=1) l++;
    for(int i=0;i<m;i++) rev[i]=rev[i>>1]>>1|(i&1)<<l-1;
    for(int i=0;i<N;i++) f1[i]=f2[i]=comp(h[i],0);
    FFT(f1,1),FFT(f2,1);
    for(int i=0;i<m;i++) f1[i]=f1[i]*f2[i];
    FFT(f1,-1);
    for(int i=1;i<N;i++) f[i]=(int)(f1[i].r+0.5);
    for(int i=1;i<N;i++)
        if(i&1) f[i]>>=1; else f[i]=(f[i]+h[i>>1])>>1;
    scanf("%d",&T);
    while(T--)
    {
        scanf("%d",&n);
        memset(tt,0,sizeof(tt));
        LL ans=0;
        for(int i=n-1;i>=0;i--)
        {
            tt[f[n]-f[i]]++;
            ans+=tt[f[i]];
        }
        printf("%lld\n",ans);
    }
    return 0;
}