UVA12879 Golf Bot FFT

传送门

题意：给出一个长度为$N$的数列$a_i$，再给出$M$个目标数字$d_i$，问目标数字中有多少个等于$a_i$或$a_i+a_j$($i,j \in [1,N]$，不要求$i,j$不相等)。$1 \leq N , M , a_i , d_i \leq 2 \times 10^5$

---

首先，等于$a_i$相当于等于$a_i+0$，那么把$0$丢进$a_i$中就变成了询问目标数字中有多少个等于$a_i$中两个数字之和。

考虑到$a_i$与$a_j$会对$a_i+a_j$产生贡献，与卷积类似，故使用$FFT$解决。与生成函数类似（好像就是生成函数），将$A_{a_i}$与$B_{a_i}$$(i \in [1 , N])$设为$1$，其余设为$0$，做$FFT$。设$FFT$结果数组为$C$，那么如果$C_{d_i}>0$，意味着$d_i$是能被$a_i$中某两个元素表示的，答案++即可。

#include<bits/stdc++.h>
#define ld long double
#define eps 1e-2
//This code is written by Itst
using namespace std;

inline int read(){
    int a = 0;
    bool f = 0;
    char c = getchar();
    while(c != EOF && !isdigit(c)){
        if(c == '-')
            f = 1;
        c = getchar();
    }
    while(c != EOF && isdigit(c)){
        a = (a << 3) + (a << 1) + (c ^ '0');
        c = getchar();
    }
    return f ? -a : a;
}

const int MAXN = 530010;
struct comp{
    ld x , y;

    comp(ld _x = 0 , ld _y = 0){
        x = _x;
        y = _y;
    }

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

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

    comp operator *(comp a){
        return comp(x * a.x - y * a.y , x * a.y + y * a.x);
    }
}A[MAXN] , B[MAXN];
int need , dir[MAXN];
const ld pi = acos(-1);

bool cmp(ld a , ld b){
    return a - eps < b && a + eps > b;
}

inline void FFT(comp* a , int type){
    for(int i = 1 ; i < need ; ++i)
        if(i < dir[i])
            swap(a[i] , a[dir[i]]);
    for(int i = 1 ; i < need ; i <<= 1){
        comp wn(cos(pi / i) , type * sin(pi / i));
        for(int j = 0 ; j < need ; j += i << 1){
            comp w(1 , 0);
            for(int k = 0 ; k < i ; ++k , w = w * wn){
                comp x = a[j + k] , y = a[i + j + k] * w;
                a[j + k] = x + y;
                a[i + j + k] = x - y;
            }
        }
    }
}

int main(){
#ifndef ONLINE_JUDGE
    freopen("12879.in" , "r" , stdin);
    freopen("12879.out" , "w" , stdout);
#endif
    while(int N = read()){
        memset(&A , 0 , sizeof(A));
        memset(&B , 0 , sizeof(B));
        int maxN = 0 , cnt = 0;
        for(int i = 1 ; i <= N ; ++i){
            int a = read();
            A[a].x = B[a].x = 1;
            maxN = max(maxN , a);
        }
        A[0].x = B[0].x = 1;
        need = 1;
        while(need <= maxN << 1)
            need <<= 1;
        for(int i = 1 ; i < need ; ++i)
            dir[i] = (dir[i >> 1] >> 1) | (i & 1 ? need >> 1 : 0);
        FFT(A , 1);
        FFT(B , 1);
        for(int i = 0 ; i < need ; ++i)
            A[i] = A[i] * B[i];
        FFT(A , -1);
        for(int M = read() ; M ; --M)
            if(!cmp(0 , A[read()].x))
                ++cnt;
        printf("%d\n" , cnt);
    }
    return 0;
}