CC Arithmetic Progressions (FFT + 分块处理)

分类： ACM_数学类2013-07-25
21:27 505人阅读 评论(0) 收藏 举报

转载请注明出处，谢谢http://blog.csdn.net/ACM_cxlove?viewmode=contents by---cxlove

题目：给出n个数，选出三个数，按下标顺序形成等差数列

http://www.codechef.com/problems/COUNTARI

如果只是形成 等差数列并不难，大概就是先求一次卷积，然后再O(n)枚举，判断2 * a[i]的种数，不过按照下标就不会了。

有种很矬的，大概就是O(n)枚举中间的数，然后 对两边分别卷积，O(n * n * lgn)。

如果能想到枚举中间的数的话，应该可以进一步想到分块处理。

如果分为K块

那么分为几种情况 ：

三个数都是在当前块中，那么可以枚举后两个数，查找第一个数，复杂度O(N/K * N/K)

两个数在当前块中，那么另外一个数可能在前面，也可能在后面，同理还是枚举两个数，查找，复杂度

O(N/K * N/K)

如果只有一个数在当前块中，那么就要对两边的数进行卷积，然后枚举当前块中的数，查询2 × a[i]。复杂度O(N * lg N)

那么总体就是O(k * (N/K * N/K + N * lg N))。

[cpp] view
plaincopy

#include <iostream>

#include <cstdio>

#include <cstring>

#include <cmath>

#include <algorithm>

using namespace std;

//FFT copy from kuangbin

const double pi = acos (-1.0);

// Complex z = a + b * i

struct Complex {

double a, b;

Complex(double _a=0.0,double _b=0.0):a(_a),b(_b){}

Complex operator + (const Complex &c) const {

return Complex(a + c.a , b + c.b);

}

Complex operator - (const Complex &c) const {

return Complex(a - c.a , b - c.b);

}

Complex operator * (const Complex &c) const {

return Complex(a * c.a - b * c.b , a * c.b + b * c.a);

}

};

//len = 2 ^ k

void change (Complex y[] , int len) {

for (int i = 1 , j = len / 2 ; i < len -1 ; i ++) {

if (i < j) swap(y[i] , y[j]);

int k = len / 2;

while (j >= k) {

j -= k;

k /= 2;

}

if(j < k) j += k;

}

}

// FFT

// len = 2 ^ k

// on = 1 DFT on = -1 IDFT

void FFT (Complex y[], int len , int on) {

change (y , len);

for (int h = 2 ; h <= len ; h <<= 1) {

Complex wn(cos (-on * 2 * pi / h), sin (-on * 2 * pi / h));

for (int j = 0 ; j < len ; j += h) {

Complex w(1 , 0);

for (int k = j ; k < j + h / 2 ; k ++) {

Complex u = y[k];

Complex t = w * y [k + h / 2];

y[k] = u + t;

y[k + h / 2] = u - t;

w = w * wn;

}

}

}

if (on == -1) {

for (int i = 0 ; i < len ; i ++) {

y[i].a /= len;

}

}

}

const int N = 100005;

typedef long long LL;

int n , a
;

int block , size;

LL num[N << 2];

int min_num = 30000 , max_num = 1;

int before
= {0}, behind
= {0} , in
= {0};

Complex x1[N << 2] ,x2[N << 2];

int main () {

#ifndef ONLINE_JUDGE

freopen("input.txt" , "r" , stdin);

#endif

scanf ("%d", &n);

for (int i = 0 ; i < n ; ++ i) {

scanf ("%d", &a[i]);

behind[a[i]] ++;

min_num = min (min_num , a[i]);

max_num = max (max_num , a[i]);

}

LL ret = 0;

block = min(n , 35);

size = (n + block - 1) / block;

for (int t = 0 ; t < block ; t ++) {

int s = t * size , e = (t + 1) * size;

if (e > n) e = n;

for (int i = s ; i < e ; i ++) {

behind[a[i]] --;

}

for (int i = s ; i < e ; i ++) {

for (int j = i + 1 ; j < e ; j ++) {

int m = 2 * a[i] - a[j];

if(m >= 1 && m <= 30000) {

// both of three in the block

ret += in[m];

// one of the number in the pre block

ret += before[m];

}

m = 2 * a[j] - a[i];

if (m >= 1 && m <= 30000) {

// one of the number in the next block

ret += behind[m];

}

}

in[a[i]] ++;

}

// pre block , current block , next block

if (t > 0 && t < block - 1) {

int l = 1;

int len = max_num + 1;

while (l < len * 2) l <<= 1;

for (int i = 0 ; i < len ; i ++) {

x1[i] = Complex (before[i] , 0);

}

for (int i = len ; i < l ; i ++) {

x1[i] = Complex (0 , 0);

}

for (int i = 0 ; i < len ; i ++) {

x2[i] = Complex (behind[i] , 0);

}

for (int i = len ; i < l ; i ++) {

x2[i] = Complex (0 , 0);

}

FFT (x1 , l , 1);

FFT (x2 , l , 1);

for (int i = 0 ; i < l ; i ++) {

x1[i] = x1[i] * x2[i];

}

FFT (x1 , l , -1);

for (int i = 0 ; i < l ; i ++) {

num[i] = (LL)(x1[i].a + 0.5);

}

for (int i = s ; i < e ; i ++) {

ret += num[a[i] << 1];

}

}

for (int i = s ; i < e ; i ++) {

in[a[i]] --;

before[a[i]] ++;

}

}

printf("%lld\n", ret);

return 0;

}