Description
给定一个长度为N的数组A[]，求有多少对i, j, k（1<=i<j<k<=N）满足A[k]-A[j]=A[j]-A[i]。
Input
第一行一个整数N（N<=10^5）。
接下来一行N个数A[i]（A[i]<=30000）。
Output
一行一个整数。
Sample Input
10
3 5 3 6 3 4 10 4 5 2
Sample Output
9
分块FFT
先分块， 暴力求出有三个在同一块，和两个在同一块的答案
三个都不在一块的FFT
#include <cmath>
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
inline int read()
{
int x=0,f=1;char ch=getchar();
while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
return x*f;
}
const int MAXN = 1e6;
const int MOD = 998244353;
long long pow_mod(long long a, long long b, long long P)
{
long long ans = 1;
while (b)
{
if (b & 1) ans = ans * a % P;
b >>= 1;
a = a * a % P;
}
return ans;
}
long long Inv, N;
int rev[MAXN];
void FFt(long long *a, int op)
{
long long wn, w, t;
for (int i = 0; i < N; i++)
if (i < rev[i])
swap(a[i], a[rev[i]]);
for (int k = 2; k <= N; k <<= 1)
{
wn = pow_mod(3, op == 1 ? (MOD - 1) / k : MOD - 1 - (MOD - 1) / k, MOD);
for (int j = 0; j < N; j += k)
{
w = 1;
for (int i = 0; i < (k >> 1); i++, w = w * wn % MOD)
{
t = a[i + j + (k >> 1)] * w % MOD;
a[i + j + (k >> 1)] = (a[i + j] - t + MOD) % MOD;
a[i + j] = (a[i + j] + t) % MOD;
}
}
}
if (op == -1)
for (int i = 0; i < N; i++)
a[i] = a[i] * Inv % MOD;
}
int Sum1[MAXN], Sum2[MAXN];
int W[MAXN];
int l[MAXN], r[MAXN];
long long A[65537], B[65537];
int main()
{
int n = read();
for (int i = 0; i < n; i++)
W[i] = read();
int len = min((int)sqrt(n) * 6, n);
int num = n / len;
if (n % len) num++;
long long ans = 0;
for (int i = 1; i <= num; i++)
{
l[i] = (i - 1) * len;
r[i] = min(l[i] + len - 1, n - 1);
}
for (int i = 1; i <= num; i++)
{
for (int j = l[i]; j <= r[i]; j++)
{
for (int k = j + 1; k <= r[i]; k++)
if (2 * W[j] - W[k] >= 0)
ans += Sum1[2 * W[j] - W[k]];
Sum1[W[j]]++;
}
}
for (int i = num; i >= 1; i--)
{
for (int j = l[i]; j <= r[i]; j++)
{
for (int k = j + 1; k <= r[i]; k++)
if (2 * W[k] - W[j] >= 0)
ans += Sum2[2 * W[k] - W[j]];
}
for (int j = l[i]; j <= r[i]; j++)
Sum2[W[j]]++;
}
N = 65536;
Inv = pow_mod(N, MOD - 2, MOD);
for (int i = 1; i < N; i++)
if (i & 1)
rev[i] = (rev[i >> 1] >> 1) | (N >> 1);
else
rev[i] = (rev[i >> 1] >> 1);
for (int i = 2; i < num; i++)
{
memset (A, 0, sizeof (A));
memset (B, 0, sizeof (B));
for (int j = 0; j < l[i]; j++) A[W[j]]++;
for (int j = r[i] + 1; j < n; j++) B[W[j]]++;
FFt(A, 1), FFt(B, 1);
for (int j = 0; j < N; j++) A[j] = A[j] * B[j];
FFt(A, -1);
for (int j = l[i]; j <= r[i]; j++) ans += A[2 * W[j]];
}
printf ("%lld\n", ans);
}