A+B Problem,Kattis - aplusb,FFT

xing_mo

于 2020-02-19 12:05:54 发布

阅读量209

点赞数

分类专栏：数学

本文链接：https://blog.csdn.net/xing_mo/article/details/104390477

版权

数学专栏收录该内容

16 篇文章 0 订阅

订阅专栏

A+B Problem,Kattis - aplusb

https://vjudge.net/problem/Kattis-aplusb/origin
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.
思路：首先因为存在负数，平移下使之都为正数，用个数组标记一下每个数出现的次数，用FFT自卷积一下就可以得出两个数相加后每个数能够得到的方案数（记得减去相同的数相加的贡献），然后枚举每个数，答案加上两个数相加得到该数的方案数即可，记得减去有1个数是0的情况

#include<bits/stdc++.h>
#define ll long long
#define ri register int
#define MAXN 100005
#define OFF 50000
using namespace std;
const double pi=acos(-1.0);
struct cpx
{
    double r, i;
    inline void operator +=(const cpx &b){ r+=b.r, i+=b.i;}
    inline cpx operator +(const cpx &b)const{ return (cpx){r+b.r, i+b.i};}
    inline cpx operator -(const cpx &b)const{ return (cpx){r-b.r, i-b.i};}
    inline cpx operator *(const cpx &b)const{ return (cpx){r*b.r-i*b.i, r*b.i+i*b.r};}
    inline cpx operator *(const double b)const{ return (cpx){r*b, i*b};}
    inline cpx operator ~()const{return (cpx){r, -i};}
}a[MAXN<<2], w[MAXN<<2];

inline void DFT_(cpx *f, int n)
{
    for(ri i=0, j=0; i<n; ++i)
    {
        if(i>j) swap(f[i], f[j]);
        for(ri k=n>>1; (j^=k)<k; k>>=1);
    }
    for(ri i=1; i<n; i<<=1) for(ri j=0; j<n; j+=i<<1)
            for(ri k=j; k<j+i; ++k)
            {
                cpx t=w[i+k-j]*f[k+i];
                f[k+i]=f[k]-t, f[k]+=t;
            }
}
inline void DFT(cpx *f, int n)
{
    if(n==1) return;
    n>>=1;
    static cpx a[MAXN<<1];
    for(ri i=0; i<n; ++i) a[i]=(cpx){f[i<<1].r, f[i<<1|1].r};
    DFT_(a, n);
    for(ri i=0; i<n; ++i)
    {
        cpx q=~a[(n-i)&(n-1)], x=(a[i]+q)*0.5, y=(a[i]-q)*(cpx){0, -0.5}, t=y*w[n+i];
        f[i]=x+t, f[n+i]=x-t;
    }
}
inline void IDFT(cpx *f, int n)
{
    if(n==1) return;
    reverse(f+1, f+n), n>>=1;
    static cpx a[MAXN<<1];
    for(ri i=0; i<n; ++i)
        a[i]=(f[i]+f[i+n])*0.5 + (f[i]-f[i+n])*(cpx){0, 0.5}*w[n+i];
    DFT_(a, n);
    double k=1.0/n;
    for(ri i=0; i<n; ++i) f[i<<1]=(cpx){a[i].r*k, 0}, f[i<<1|1]=(cpx){a[i].i*k, 0};
}

int n,x[MAXN],len1,tmp,b[MAXN],u[MAXN];
ll ans[MAXN << 2];
int main()
{
    while(~scanf("%d",&n))//输入两个多项式的长度及模数,注意这里的长度比实际长度小1
    {
        len1 = 0;
        memset(x,0,sizeof(x));
        for(int i = 1;i <= n;++i)
        {
            scanf("%d",&u[i]);
            ++x[u[i] + OFF];
            len1 = max(len1,u[i]+OFF);
        }
        ++len1;
        int len = 1;
        while(len <(len1<<1)) len<<=1;
        for(int i = 0;i < len1;i++)
            a[i] = (cpx){(double)x[i],0};
        for(int i = len1;i < len;i++)
            a[i] = (cpx){0,0};
        for(ri i=1; i<len; i<<=1)
        {
            w[i]=(cpx){1, 0};
            for(ri j=1; j<i; ++j)
                w[i+j]=((j&31)==1?(cpx){cos(pi*j/i), sin(pi*j/i)}:w[i+j-1]*w[i+1]);
        }

        DFT(a,len);
        for(int i = 0;i < len;++i)
            a[i] = a[i]*a[i];
        IDFT(a,len);

        memset(ans,0,sizeof(ans));
        ll res = 0;
        for(int i = 0;i < len;++i)
            ans[i] = (ll)(0.5+a[i].r);
        for(int i = 1;i <= n;++i)
            --ans[u[i]*2+2*OFF];
        for(int i = 1;i <= n;++i)
        {
            res += ans[u[i]+2*OFF];
            res -= 2*(x[OFF]-(u[i]==0));
        }
        printf("%lld\n",res);
    }
    return 0;
}