CodeChef - COUNTARI FTT+分塊
Arithmetic Progressions
Given N integers A1, A2, …. AN, Dexter wants to know how many ways he can choose three numbers such that they are three consecutive terms of an arithmetic progression.
Meaning that, how many triplets (i, j, k) are there such that 1 ≤ i < j < k ≤ Nand Aj
So the triplets (2, 5, 8), (10, 8, 6), (3, 3, 3) are valid as they are three consecutive terms of an arithmetic
progression. But the triplets (2, 5, 7), (10, 6, 8) are not.
Input
First line of the input contains an integer N (3 ≤ N ≤ 100000). Then the following line contains N space separated integers A1
Output
Output the number of ways to choose a triplet such that they are three consecutive terms of an arithmetic progression.
Example
Input: 10 3 5 3 6 3 4 10 4 5 2 Output: 9
Explanation
The followings are all 9 ways to choose a triplet
1 : (i, j, k) = (1, 3, 5), (Ai, Aj, Ak) = (3, 3, 3) 2 : (i, j, k) = (1, 6, 9), (Ai, Aj, Ak) = (3, 4, 5) 3 : (i, j, k) = (1, 8, 9), (Ai, Aj, Ak) = (3, 4, 5) 4 : (i, j, k) = (3, 6, 9), (Ai, Aj, Ak) = (3, 4, 5) 5 : (i, j, k) = (3, 8, 9), (Ai, Aj, Ak) = (3, 4, 5) 6 : (i, j, k) = (4, 6, 10), (Ai, Aj, Ak) = (6, 4, 2) 7 : (i, j, k) = (4, 8, 10), (Ai, Aj, Ak) = (6, 4, 2) 8 : (i, j, k) = (5, 6, 9), (Ai, Aj, Ak) = (3, 4, 5) 9 : (i, j, k) = (5, 8, 9), (Ai, Aj, Ak) = (3, 4, 5)
題解:
考慮分塊,分成block塊
假設三個點都在同一塊,那麽我們就在一塊內暴力,復雜度block * ( n/block) * (n/block)
假設其中兩個點在同一塊,那麽枚舉其中一塊的兩個點算答案,block * n/block * n/block
· 假設三個點都不在同一塊,枚舉中間點屬於的那一塊 剩下左邊和右邊進行 FFT, 復雜度block * (n*logn)
#include<bits/stdc++.h> using namespace std; #pragma comment(linker, "/STACK:102400000,102400000") #define ls i<<1 #define rs ls | 1 #define mid ((ll+rr)>>1) #define pii pair<int,int> #define MP make_pair typedef long long LL; typedef unsigned long long ULL; const long long INF = 1e18+1LL; const double pi = acos(-1.0); const int N = 3e5+20, M = 1e6+10, mod = 1e9+7,inf = 2e9; struct Complex { double r , i ; Complex () {} Complex ( double r , double i ) : r ( r ) , i ( i ) {} Complex operator + ( const Complex& t ) const { return Complex ( r + t.r , i + t.i ) ; } Complex operator - ( const Complex& t ) const { return Complex ( r - t.r , i - t.i ) ; } Complex operator * ( const Complex& t ) const { return Complex ( r * t.r - i * t.i , r * t.i + i * t.r ) ; } } ; void FFT ( Complex y[] , int n , int rev ) { for ( int i = 1 , j , t , k ; i < n ; ++ i ) { for ( j = 0 , t = i , k = n >> 1 ; k ; k >>= 1 , t >>= 1 ) j = j << 1 | t & 1 ; if ( i < j ) swap ( y[i] , y[j] ) ; } for ( int s = 2 , ds = 1 ; s <= n ; ds = s , s <<= 1 ) { Complex wn = Complex ( cos ( rev * 2 * pi / s ) , sin ( rev * 2 * pi / s ) ) , w ( 1 , 0 ) , t ; for ( int k = 0 ; k < ds ; ++ k , w = w * wn ) { for ( int i = k ; i < n ; i += s ) { y[i + ds] = y[i] - ( t = w * y[i + ds] ) ; y[i] = y[i] + t ; } } } if ( rev == -1 ) for ( int i = 0 ; i < n ; ++ i ) y[i].r /= n ; } Complex s[N],t[N]; LL cnt[502][30005]; int a[N]; int n,block,pos[N]; LL vis[N]; int main() { while(scanf("%d",&n)!=EOF) { block = 1500; for(int i = 1; i <= n; ++i) pos[i] = (i-1)/block + 1; int mx = -1; for(int i = 0; i <= pos[n]; ++i) for(int j = 1; j <= 30000; ++j) cnt[i][j] = 0; for(int i = 1; i <= n; ++i) { scanf("%d",&a[i]); mx = max(mx,a[i]); cnt[pos[i]][a[i]]++; } for(int i = 1; i <= mx; ++i) { for(int j = 1; j <= pos[n]; ++j) { cnt[j][i] += cnt[j-1][i]; } } int len = 1; while(len <= 2*mx) len<<=1; LL ans = 0; for(int k = 1; k <= pos[n]; ++k) { for(int i = (k-1)*block + 1; i <= min(k*block,n); ++i) { for(int j = i + 1; j <= min(k*block,n); ++j) { if(2*a[i] - a[j] >= 1 && 2*a[i] - a[j] <= mx) ans += cnt[k-1][2*a[i] - a[j]] + vis[2*a[i]-a[j]]; if(2*a[j] - a[i] >= 1 && 2*a[j] - a[i] <= mx) ans += cnt[pos[n]][2*a[j] - a[i]] - cnt[k][2*a[j] - a[i]]; } vis[a[i]] += 1; } for(int i = (k-1)*block + 1; i <= min(k*block,n); ++i) { vis[a[i]] = 0; } for(int j = 0; j <= mx; ++j) s[j] = Complex(cnt[k-1][j],0); for(int j = mx+1; j < len; ++j) s[j] = Complex(0,0); for(int j = 0; j <= mx; ++j) t[j] = Complex(cnt[pos[n]][j] - cnt[k][j] , 0); for(int j = mx+1; j < len; ++j) t[j] = Complex(0,0); FFT(s,len,1);FFT(t,len,1); for(int j = 0; j < len; ++j) s[j] = s[j] * t[j]; FFT(s,len,-1); for(int j = 1; j <= mx; ++j) { LL tmp = (LL)(s[2*j].r + 0.5); ans += tmp*(cnt[k][j] - cnt[k-1][j]); } } printf("%lld\n",ans); } return 0; }
CodeChef - COUNTARI FTT+分塊