BZOJ-3509 母函数+分块+暴力+FFT
题目描述
给定一个长度为N的数组A[],求有多少对i, j, k(1<=i<j<k<=N)满足A[k]-A[j]=A[j]-A[i]。
输入格式
第一行一个整数N(N<=10^5)。
接下来一行N个数A[i](A[i]<=30000)。
输出格式
一行一个整数。
样例输入
10 3 5 3 6 3 4 10 4 5 2
样例输出
9
这题网上题解很多
我的第一个想法是枚举 j 然后左边的数 和 右边的数 FFT处理一下
这样的复杂度是 n^2*log(n) 复杂度明显不行 所以需要进一步的优化
网上题解都是用分块进行优化
先在每一个块里面找一下是否存在满足条件的
然后再找块外面的满足条件的
这题BZOJ 开了40S 跑的了的
我FFT是直接套板子 当作黑盒使用 所有代码有点鬼畜
#include <cstdio> #include <cstring> #include <queue> #include <cmath> #include <algorithm> #include <set> #include <iostream> #include <map> #include <stack> #include <string> #include <vector> #define pi acos(-1.0) #define eps 1e-9 #define fi first #define se second #define rtl rt<<1 #define rtr rt<<1|1 #define bug printf("******\n") #define mem(a,b) memset(a,b,sizeof(a)) #define name2str(x) #x #define fuck(x) cout<<#x" = "<<x<<endl #define f(a) a*a #define sf(n) scanf("%d", &n) #define sff(a,b) scanf("%d %d", &a, &b) #define sfff(a,b,c) scanf("%d %d %d", &a, &b, &c) #define sffff(a,b,c,d) scanf("%d %d %d %d", &a, &b, &c, &d) #define pf printf #define FRE(i,a,b) for(i = a; i <= b; i++) #define FREE(i,a,b) for(i = a; i >= b; i--) #define FRL(i,a,b) for(i = a; i < b; i++)+ #define FRLL(i,a,b) for(i = a; i > b; i--) #define FIN freopen("data.txt","r",stdin) #define gcd(a,b) __gcd(a,b) #define lowbit(x) x&-x #define rep(i,a,b) for(int i=a;i<b;++i) #define per(i,a,b) for(int i=a-1;i>=b;--i) using namespace std; typedef long long LL; typedef unsigned long long ULL; const int maxn = 3e5 + 7; const int maxm = maxn * 4; int n, m, a[maxn], b[maxn]; int len, res[maxm], mx; //开大4倍 struct cpx { long double r, i; cpx ( long double r = 0, long double i = 0 ) : r ( r ), i ( i ) {}; cpx operator+ ( const cpx &b ) { return cpx ( r + b.r, i + b.i ); } cpx operator- ( const cpx &b ) { return cpx ( r - b.r, i - b.i ); } cpx operator* ( const cpx &b ) { return cpx ( r * b.r - i * b.i, i * b.r + r * b.i ); } } va[maxm], vb[maxm]; void rader ( cpx F[], int len ) { //len = 2^M,reverse F[i] with F[j] j为i二进制反转 int j = len >> 1; for ( int i = 1; i < len - 1; ++i ) { if ( i < j ) swap ( F[i], F[j] ); // reverse int k = len >> 1; while ( j >= k ) j -= k, k >>= 1; if ( j < k ) j += k; } } void FFT ( cpx F[], int len, int t ) { rader ( F, len ); for ( int h = 2; h <= len; h <<= 1 ) { cpx wn ( cos ( -t * 2 * pi / h ), sin ( -t * 2 * pi / h ) ); for ( int j = 0; j < len; j += h ) { cpx E ( 1, 0 ); //旋转因子 for ( int k = j; k < j + h / 2; ++k ) { cpx u = F[k]; cpx v = E * F[k + h / 2]; F[k] = u + v; F[k + h / 2] = u - v; E = E * wn; } } } if ( t == -1 ) //IDFT for ( int i = 0; i < len; ++i ) F[i].r /= len; } void Conv ( cpx a[], cpx b[], int len ) { //求卷积 FFT ( a, len, 1 ); FFT ( b, len, 1 ); for ( int i = 0; i < len; ++i ) a[i] = a[i] * b[i]; FFT ( a, len, -1 ); } void gao () { len = 1; mx = n + m; while ( len <= mx ) len <<= 1; //mx为卷积后最大下标 for ( int i = 0; i < len; i++ ) va[i].r = va[i].i = vb[i].r = vb[i].i = 0; for ( int i = 0; i < n; i++ ) va[i].r = a[i]; //根据题目要求改写 for ( int i = 0; i < m; i++ ) vb[i].r = b[i]; //根据题目要求改写 Conv ( va, vb, len ); for ( int i = 0; i < len; ++i ) res[i] += ( LL ) floor ( va[i].r + 0.5 ); } int L[maxn], R[maxn], p[maxn]; int main() { //FIN; int num; sf ( num ); m = -1; for ( int i = 0; i < num ; i++ ) { sf ( p[i] ); R[p[i]]++; m = max ( m, p[i] ); } m++; n = m; int sz = min ( num, 8 * ( int ) sqrt ( num ) ); LL ans = 0; for ( int i = 0 ; i < num ; i += sz ) { int r = min ( num, i + sz ) - 1; for ( int j = i ; j <= r ; j++ ) R[p[j]]--; for ( int j = i, x ; j <= r ; j++ ) { for ( int k = j + 1 ; k <= r ; k++ ) { x = 2 * p[j] - p[k]; if ( x >= 0 ) ans += 1LL * L[x]; x = 2 * p[k] - p[j]; if ( x >= 0 ) ans += 1LL * R[x]; } L[p[j]]++; } } for ( int i = 0 ; i < num ; i += sz ) { int r = min ( num, i + sz ) - 1; mem ( a, 0 ), mem ( b, 0 ), mem ( res, 0 ); for ( int j = 0 ; j < i ; j++ ) a[p[j]]++; for ( int j = r + 1 ; j < num ; j++ ) b[p[j]]++; gao(); for ( int j = i ; j <= r ; j++ ) ans += 1LL * ( res[2 * p[j]] ); } printf ( "%lld\n", ans ); return 0; }