2017北理校赛G题 人民的名义（FFT）

 

把1、2和3、4的藏钱地方的钱的和的可能枚举出来，原始的O(n^2)会tle，用fft。

然后暴力枚举query的值，拆成两半，把两部分的和乘起来。

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;
const double PI = acos(-1.0);
typedef struct Complex {
    double r,i;
    Complex(double _r = 0.0,double _i = 0.0) {
        r = _r; i = _i;
    }
    Complex operator +(const Complex &b) {
        return Complex(r+b.r,i+b.i);
    }
    Complex operator -(const Complex &b) {
        return Complex(r-b.r,i-b.i);
    }
    Complex operator *(const Complex &b) {
        return Complex(r*b.r-i*b.i,r*b.i+i*b.r);
    }
}Complex;
void change(Complex y[],LL len) {
    LL 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 /= 2;
        }
        if(j < k) j += k;
    }
}

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

const LL maxn = 69200;
LL cnt[5][maxn];
Complex f[5][maxn];
Complex r[3][maxn];
LL rx[3][maxn];
LL n, len, mx;

void gao(LL i1, LL i2, LL id) {
    for(LL i = 0; i < len; i++) f[i1][i] = Complex(0, 0);
    for(LL i = 0; i < mx; i++) f[i1][i] = Complex(cnt[i1][i], 0);
    fft(f[i1], len, 1);
    for(LL i = 0; i < len; i++) f[i2][i] = Complex(0, 0);
    for(LL i = 0; i < mx; i++) f[i2][i] = Complex(cnt[i2][i], 0);
    fft(f[i2], len, 1);
    for(LL i = 0; i < len; i++) r[id][i] = f[i1][i] * f[i2][i];
    fft(r[id], len, -1);
    for(LL i = 0; i < len; i++) rx[id][i] = (LL)(r[id][i].r + 0.5);
}

int main() {
    // freopen("in", "r", stdin);
    int T, q;
    LL x;
    scanf("%d", &T);
    while(T--) {
        memset(cnt, 0, sizeof(cnt));
        scanf("%lld",&n);
        mx = 0;
        for(LL i = 1; i <= 4; i++) {
            for(LL j = 1; j <= n; j++) {
                scanf("%lld", &x);
                mx = max(mx, x);
                cnt[i][x]++;
            }
        }
        mx++; len = 1;
        while(len < 2 * mx) len <<= 1;
        gao(1, 2, 0); gao(3, 4, 1);
        scanf("%d", &q);
        while(q--) {
            LL val;
            scanf("%lld", &val);
            LL ret = 0;
            for(LL i = 0; i <= val; i++) {
                LL l = i, r = val - i;
                ret += rx[0][l] * rx[1][r];
                // cout << l << "->" << rx[0][l] << " " << r << "->" << rx[1][r] << endl;
            }
            printf("%lld
", ret);
        }
    }
    return 0;
}