【FFT】CODECHEF COUNTARI
通道：https://www.codechef.com/problems/COUNTARI
题意：给出一个数列num[1..n], 每个数都是不超过30000的正整数, 现在求有多少个三元组(i, j, k)满足 1 <= i < j < k <= n使得 A[i], A[j], A[k]成等差数列
思路：
刚开始想到判断A[i] + A[k] == 2*A[j]可以用FFT处理出任意两个数的和为T的有多少种, 然而发现这样难处理i, j, k的顺序
可以用分块，考虑三个数的位置
1.如果三个数在同一块中, 直接对每一块进行枚举后两个数的位置然后查询当前块中需要成为等差数列需要的那个数有多少个, 时间复杂度是O(N/K * N/K *K)
2.如果三个数中有两个在同一块中, 枚举那两个数, 另外一个数可能在前面的块中也可能在后面的块中, 同上直接查询即可, 时间复杂度是 O(N/ K * N/K *K)
3.如果三个数在不同块中, 枚举中间那个数的在这当前块中的位置N/K种, 然后对于此块前后的部分作FFT得到前后各取一个数和是T有多少种, 对于N/K次直接查询即可，FFT一共进行K次, 时间复杂度是O(K*MlogM) 综上总体时间复杂度是O(N*N/K + K*MlogM),
代码：
1 #include <cstdio> 2 #include <cmath> 3 #include <cstring> 4 #include <vector> 5 #include <algorithm> 6 7 using namespace std; 8 9 typedef long long ll; 10 11 const int K = 57; 12 const int MAX_N = 100007; 13 const int MAX_M = 100007; 14 const double PI = acos(-1.0); 15 16 struct Complex { 17 double r, i; 18 Complex(double _r, double _i) { 19 r = _r; 20 i = _i; 21 } 22 Complex operator + (const Complex &c) { 23 return Complex(c.r + r, c.i + i); 24 } 25 Complex operator - (const Complex &c) { 26 return Complex(r - c.r, i - c.i); 27 } 28 Complex operator * (const Complex &c) { 29 return Complex(c.r * r - c.i * i, c.r * i + c.i * r); 30 } 31 Complex operator / (const int &c) { 32 return Complex(r / c, i / c); 33 } 34 Complex(){} 35 }; 36 namespace FFT { 37 int rev(int id, int len) { 38 int ret = 0; 39 for(int i = 0; (1 << i) < len; ++i) { 40 ret <<= 1; 41 if(id & (1 << i)) ret |= 1; 42 } 43 return ret; 44 } 45 Complex A[MAX_M << 3]; 46 void FFT(Complex *a, int len, int DFT) { 47 for(int i = 0; i < len; ++i) A[rev(i, len)] = a[i]; 48 for(int s = 1; (1 << s) <= len; ++s) { 49 int m = (1 << s); 50 Complex wm = Complex(cos(PI * DFT * 2 / m), sin(PI * DFT * 2 / m)); 51 for(int k = 0; k < len; k += m) { 52 Complex w = Complex(1, 0); 53 for(int j = 0; j < (m >> 1); j++) { 54 Complex t = w * A[k + j + (m >> 1)]; 55 Complex u = A[k + j]; 56 A[k + j] = u + t; 57 A[k + j + (m >> 1)] = u - t; 58 w = w * wm; 59 } 60 } 61 } 62 if(DFT == -1) for(int i = 0; i < len; ++i) A[i] = A[i] / len; 63 for(int i = 0; i < len; i++) a[i] = A[i]; 64 } 65 }; 66 67 int n; 68 int num[MAX_N]; 69 Complex L[MAX_M << 3], R[MAX_M << 3]; 70 bool vis[MAX_M]; 71 int before[MAX_M], behind[MAX_M], now[MAX_M]; 72 73 int main() { 74 while(1 == scanf("%d", &n)) { 75 int up = 1; 76 memset(before, 0, sizeof before); 77 memset(behind, 0, sizeof behind); 78 memset(vis, false, sizeof vis); 79 memset(now, 0, sizeof now); 80 for(int i = 1; i <= n; ++i) 81 scanf("%d", num + i), behind[num[i]]++, vis[num[i]] = true, up = max(up, num[i]); 82 int block = min(K, n), size = (n + block - 1) / block; 83 int len = 1; 84 while(len <= up) len <<= 1; len <<= 1; 85 ll ans = 0; 86 for(int pos = 1; pos <= block; ++pos) { 87 int s = size * (pos - 1) + 1, e = min(n, size * pos); 88 for(int i = s; i <= e; ++i) --behind[num[i]]; 89 for(int i = s; i <= e; ++i) { 90 for(int j = i + 1; j <= e; ++j) { 91 int a = 2 * num[i] - num[j]; 92 if(a >= 1 && a <= 30000 && vis[a]) 93 ans += now[a] + before[a]; 94 int c = 2 * num[j] - num[i]; 95 if(c >= 1 && c <= 30000 && vis[c]) 96 ans += behind[c]; 97 } 98 now[num[i]]++; 99 } 100 for(int i = 0; i <= up; ++i) { 101 L[i] = Complex(before[i], 0); 102 R[i] = Complex(behind[i], 0); 103 } 104 for(int i = up + 1; i < len; ++i) 105 L[i] = R[i] = Complex(0, 0); 106 FFT::FFT(L, len, 1); FFT::FFT(R, len, 1); 107 for(int i = 0; i < len; ++i) 108 L[i] = L[i] * R[i]; 109 FFT::FFT(L, len, -1); 110 for(int j = s; j <= e; ++j) 111 ans += (ll)(L[2 * num[j]].r + 0.5); 112 for(int j = s; j <= e; j++) ++before[num[j]], --now[num[j]]; 113 } 114 printf("%lld\n", ans); 115 } 116 return 0; 117 }
