BZOJ3509: [CodeChef] COUNTARI

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#

1
10
2
3 5 3 6 3 4 10 4 5 2

Sample Output#

1
9

分块FFT
先分块，暴力求出有三个在同一块，和两个在同一块的答案
三个都不在一块的FFT

1
#include <cmath>
2
#include <cstdio>
3
#include <cstring>
4
#include <algorithm>
5
using namespace std;
6
inline int read()
7
{
8
    int x=0,f=1;char ch=getchar();
9
    while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
10
    while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
11
    return x*f;
12
}
13
const int MAXN = 1e6;
14
const int MOD = 998244353;
15
long long pow_mod(long long a, long long b, long long P)
16
{
17
    long long ans = 1;
18
    while (b)
19
    {
20
        if (b & 1) ans = ans * a % P;
21
        b >>= 1;
22
        a = a * a % P;
23
    }
24
    return ans;
25
}
26
long long Inv, N;
27
int rev[MAXN];
28
void FFt(long long *a, int op)
29
{
30
    long long wn, w, t;
31
    for (int i = 0; i < N; i++)
32
        if (i < rev[i])
33
            swap(a[i], a[rev[i]]);
34
    for (int k = 2; k <= N; k <<= 1)
35
    {
36
        wn = pow_mod(3, op == 1 ? (MOD - 1) / k : MOD - 1 - (MOD - 1) / k, MOD);
37
        for (int j = 0; j < N; j += k)
38
        {
39
            w = 1;
40
            for (int i = 0; i < (k >> 1); i++, w = w * wn % MOD)
41
            {
42
                t = a[i + j + (k >> 1)] * w % MOD;
43
                a[i + j + (k >> 1)] = (a[i + j] - t + MOD) % MOD;
44
                a[i + j] = (a[i + j] + t) % MOD;
45
            }
46
        }
47
    }
48
    if (op == -1)
49
        for (int i = 0; i < N; i++)
50
            a[i] = a[i] * Inv % MOD;
51
}
52
int Sum1[MAXN], Sum2[MAXN];
53
int W[MAXN];
54
int l[MAXN], r[MAXN];
55
long long A[65537], B[65537];
56
int main()
57
{
58
    int n = read();
59
    for (int i = 0; i < n; i++)
60
        W[i] = read();
61
    int len = min((int)sqrt(n) * 6, n);
62
    int num = n / len;
63
    if (n % len) num++;
64
    long long ans = 0;
65
    for (int i = 1; i <= num; i++)
66
    {
67
        l[i] = (i - 1) * len;
68
        r[i] = min(l[i] + len - 1, n - 1);
69
    }
70
    for (int i = 1; i <= num; i++)
71
    {
72
        for (int j = l[i]; j <= r[i]; j++)
73
        {
74
            for (int k = j + 1; k <= r[i]; k++)
75
                if (2 * W[j] - W[k] >= 0)
76
                    ans += Sum1[2 * W[j] - W[k]];
77
            Sum1[W[j]]++;
78
        }
79
    }
80
    for (int i = num; i >= 1; i--)
81
    {
82
        for (int j = l[i]; j <= r[i]; j++)
83
        {
84
            for (int k = j + 1; k <= r[i]; k++)
85
                if (2 * W[k] - W[j] >= 0)
86
                    ans += Sum2[2 * W[k] - W[j]];
87
        }
88
        for (int j = l[i]; j <= r[i]; j++)
89
            Sum2[W[j]]++;
90
    }
91
    N = 65536;
92
    Inv = pow_mod(N, MOD - 2, MOD);
93
    for (int i = 1; i < N; i++)
94
        if (i & 1)
95
            rev[i] = (rev[i >> 1] >> 1) | (N >> 1);
96
        else
97
            rev[i] = (rev[i >> 1] >> 1);
98
    for (int i = 2; i < num; i++)
99
    {
100
        memset (A, 0, sizeof (A));
101
        memset (B, 0, sizeof (B));
102
        for (int j = 0; j < l[i]; j++) A[W[j]]++;
103
        for (int j = r[i] + 1; j < n; j++) B[W[j]]++;
104
        FFt(A, 1), FFt(B, 1);
105
        for (int j = 0; j < N; j++) A[j] = A[j] * B[j];
106
        FFt(A, -1);
107
        for (int j = l[i]; j <= r[i]; j++) ans += A[2 * W[j]];
108
    }
109
    printf ("%lld\n", ans);
110
}