2018-01-22 A+B Problem 16年香港网络赛 FFT

鉴于这个题目所在的OJ不太好找, 故这里直接附上传送门:

https://open.kattis.com/problems/aplusb

题目描述:
Given N integers in the range [−50000,50000] , how many ways are there to pick three integers ai , aj , ak , such that i , j , k are pairwise distinct and ai+aj=ak ? Two ways are different if their ordered triples (i,j,k) of indices are different.

输入:
The first line of input consists of a single integer N ( 1≤N≤200000 ). The next line consists of N space-separated integers a1,a2,…,aN .

输出:
Output an integer representing the number of ways.

样例输入:
4
1 2 3 4
6
1 1 3 3 4 6

样例输出:
4
10

题目的意思很简单, 就是给你一排数, 然后给定他们的下标, 问你有几种 ai , aj , ak 可以满足 ai+aj=ak,

其中 然后, 值得注意的是, aj+ai=ak会被认为是一种不同的加的式子.

老套路, 先是用FFT处理出一组生成函数的系数的数组, 例如res[4]=2表示两个数字加起来等于4的方案数为2.

当然, 这个也是要去重的, 不过比上一篇博文简单一些, 主要是因为 "aj+ai=ak会被认为是一种不同的加的式子" .

因此只需要把选了两次一样的数的重复去掉, 然后不需要把所有的次数都除以2.

然后遍历所有的数字, 其对应数值可以在res[]数组里面找到对应的方案数, 加起来就行,取为ans.

最后还有一个去重. 比如当这个数列有0的时候.

行啦, 比如有n个数字, 里面有z个0.

当我遍历到一个非0的数字的时候, 对应的res[]里面会包含2*z种方案---2*z种这个数加0的方案, z要乘以2因为提到"aj+ai=ak会被认为是一种不同的加的式子" 有n-z个非零的数字, 那么意味着错误的方案有 2*n*(n-z)

当我遍历到一个0的时候, 对应的res[]里面会有2*(z-1)种方案---2*(z-1)这个数加其他0还是本身的方案, 是2*(z-1)而不是(z-1)的原因同上, 有个z个0, 那么错误的方案数还有2*z*(z-1)

所以 ans再减去2*n*(n-z), 再减去2*z*(z-1), 就是结果了

代码如下:

#include <iostream>
#include <cstdlib>
#include <cstring>
#include <cstdio>
#include <cmath>
#include <algorithm>
using namespace std;
const double PI=acos(-1);
struct complex
{
    double a,b;
    complex(double aa=0.0,double bb=0.0)
    {
        a=aa;
        b=bb;
    }
    complex operator +(const complex &e)
    {
        return complex(a+e.a,b+e.b);
    }
    complex operator -(const complex &e)
    {
        return complex(a-e.a,b-e.b);
    }
    complex operator *(const complex &e)
    {
        return complex(a*e.a-b*e.b,a*e.b+b*e.a);
    }
};

void change(complex *y,int len)
{
    int i,j,k;
    for(i=1,j=len/2;i<len-1; i++)
    {
        if(i<j)
            swap(y[i],y[j]);
        k=len/2;
        while(j>=k)
        {
            j-=k;
            k>>=1;
        }
        if(j<k)
            j+=k;
    }
}
void fft(complex *y,int len,int on)
{
    change(y, len);
    for(int h=2; h<=len; h<<=1)
    {
        complex wn(cos(-on*2*PI/h),sin(-on*2*PI/h));
        for(int j=0; j<len; j+=h)
        {
            complex w(1, 0);
            for(int k=j;k<j+h/2;k++)
            {
                complex u=y[k],t=w*y[k+h/2];
                y[k]=u+t;
                y[k+h/2]=u-t;
                w=w*wn;
            }
        }
    }
    if(on == -1)
        for(int i=0; i<len; i++)
            y[i].a /= len;
}

const int k=50000;
const int maxx=800040;
int data[maxx];
long long int times[maxx],res[maxx];
complex a[maxx],b[maxx];
int len_a,len_b;

int main()
{
    int n,i;
    long long int ans;
    while(scanf("%d",&n)!=EOF)
    {
        memset(data,0,sizeof(data));
        memset(res,0,sizeof(res));
        memset(times,0,sizeof(times));
        memset(a,0,sizeof(a));
        memset(b,0,sizeof(b));
        for(i=0;i<n;i++)
        {
            scanf("%d",&data[i]);
            times[data[i]+k]++;
        }
        sort(data,data+n);
        len_a=data[n-1]+k+1;
        len_b=1;
        while(len_b<(len_a*2))
            len_b=len_b<<1;
        for(i=0;i<len_b;i++)
            a[i].a=times[i];
        fft(a,len_b,1);
        for(i=0;i<len_b;i++)
            b[i]=a[i]*a[i];
        fft(b,len_b,-1);
        for(i=0;i<len_b;i++)
            res[i]=(long long int)(b[i].a+0.5);
        for(i=0;i<n;i++)
            res[2*data[i]+2*k]--;
        ans=0;
        for(i=0;i<n;i++)
            ans+=res[data[i]+2*k];
        ans-=2*times[k]*(n-times[k])+times[k]*(times[k]-1)*2;
        printf("%lld\n",ans);
    }
    return 0;
}