转载请注明出处,谢谢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))。
- #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[N];
- int block , size;
- LL num[N << 2];
- int min_num = 30000 , max_num = 1;
- int before[N] = {0}, behind[N] = {0} , in[N] = {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;
- }