BZOJ 3509 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

10
3 5 3 6 3 4 10 4 5 2

Sample Output

9

Solution

首先我们转化一下条件,显然我们要求一个 \( 2*A_j=A_k+A_i \) ,
也就是找长度为3的有序等差数列.
考虑生成函数 ,
对于两个序列,我们对他FFT得到的就是每一项与另一个序列每一项乘积 ,
然后说分块 ,
分成 \(sz\) 块 .
那么我们对于块内转移可以暴力解决,\(j\)在块内,\(i\)、\(k\)在块外的情况FFT解决 .

代码:

#include <cstdio>
#include <cmath>
#include <cstring>
#include <iostream>
using namespace std;
const double pi = acos(-1);
typedef long long ll;
struct Complex {
    double x, y;
    Complex() {}
    Complex(double tx,double ty) {x = tx;y = ty;}
    Complex operator + (Complex b) {return Complex(x + b.x,y + b.y);}
    Complex operator - (Complex b) {return Complex(x - b.x,y - b.y);}
    Complex operator * (Complex b) {return Complex(x * b.x - y * b.y,x * b.y + y * b.x);}
    Complex operator *= (Complex b) {*this = Complex(x * b.x - y * b.y,x * b.y + y * b.x);return *this;}
};
int getint(){
    int r=0;char c;
    for (c=getchar();c<'0'||c >'9';c=getchar());
    for (;'0'<=c&&c<='9'; c=getchar()) r = r * 10 - '0' + c;
    return r;
}
const int maxm = 700010;
Complex a[maxm], b[maxm];
int n, m, L, R[maxm], num[maxm], r[maxm], st[maxm], ed[maxm], l[maxm];
int nn;
void fft(Complex *a, int f) {
    for (int i = 0; i < n; ++i) if (i < R[i]) swap(a[i], a[R[i]]);
    for (int i = 1; i < n; i <<= 1) {
        Complex wn(cos(pi / i), f * sin(pi / i));
        for (int j = 0; j < n; j += (i << 1)) {
            Complex w(1, 0);
            for (int k = 0; k < i; ++k) {
                Complex x = a[j + k], y = w * a[j + k + i];
                a[j + k] = x + y; a[j + k + i] = x - y;
                w *= wn;
            }
        }
    }
    if (f == -1) for (int i = 0; i < n; ++i) a[i].x /= n;
}
int mx = 0, size = 2000;
ll ans = 0;
int main() {
    scanf("%d", &nn);
    for (int i = 1; i <= nn; ++i) {
        scanf("%d", num+i);
        ++ r[num[i]];
        mx = max(mx, num[i]);
    }
    ++mx;
    L = 0; mx = mx * 2 - 1; for (n = 1; n <= mx; n <<= 1) ++L;
    R[0] = 0; for (int i = 1; i < n; ++i) R[i] = (R[i >> 1] >> 1) | ((i & 1) << (L-1));
    m = (nn-1)/size+1;
    for (int i = 1; i <= m; ++i) {
        st[i] = ed[i-1]+1;
        ed[i] = i*size;
    }
    ed[m] = nn;
    for (int i = 1; i <= m; ++i) {
        for (int j = st[i]; j<=ed[i];++j) r[num[j]]--;
        for (int j = 0; j < n; ++j) a[j] = Complex(l[j],0);
        for (int j = 0; j < n; ++j) b[j] = Complex(r[j],0);
        fft(a,1);
        fft(b,1);
        for(int j=0; j<n; ++j) a[j] *= b[j];
        fft(a,-1);
        for(int j = st[i];j<=ed[i];++j)
            ans += ll(a[num[j]*2].x+0.5);
        for(int j = st[i];j<=ed[i];++j) {
            for(int k=st[i]; k<j; ++k)
                if(num[j]*2-num[k] >= 0)
                    ans += ll(r[num[j]*2-num[k]]);
            for (int k = j+1; k<=ed[i]; ++k)
                if(num[j]*2-num[k] >= 0)
                    ans+= ll(l[num[j]*2-num[k]]);
            l[num[j]]++;
        }
    }
    printf("%lld\n", ans);
    return 0;
}

发表评论

电子邮件地址不会被公开。 必填项已用*标注

此站点使用Akismet来减少垃圾评论。了解我们如何处理您的评论数据