题目链接:https://www.codechef.com/problems/COUNTARI
题目大意:给出一个长度为n的数组,要求从中选出三个数a[i],a[j],a[k],满足i<j<k,且a[j]-a[i]=a[k]-a[j]。问能选出多少个这样的三元组。
题目思路:由于取出的三元组为等差数列,如果没有i < j < k的限制,我们就可以直接枚举中间那个数 a[j] 的值,借助FFT求出有多少二元组的和为2*a[j]的数量即可。但现在加入了i < j < k之后,我们就无法直接用枚举中间数的方法了,因为顺序无法保证满足i < j < k。在参考了大佬的博客之后,学会了如何用分块来解决这个问题。
我们先将这n个数分为block块,每一块的数的个数为n/block个(也可能不足)。接着对可能为答案的三元组进行分类讨论。
第一类为a[i]处于前面的块内,a[j]和a[k]处于当前的块内;
第二类为a[i]、a[j]、a[k]都处于当前的块内;
第三类为a[i]和a[j]处于当前块内,a[k]处于后面的块内;
第四类为a[i]处于前面的块内,a[j]处于当前的块内,a[k]处于后面的块内。
我们用Front[i] 表示前面的块内值为 i 的数量,Now[i] 表示当前的块内值为 i 的数量,Behind[i] 表示后面的块内值为 i 的数量。
对于前三种情况,我们可以枚举其中的两个值,再求第三个值的数量,复杂度为O(n/block * n/block),对于第四种情况来说,我们就可以将前面的块的值和后面的块的值做一次FFT,再枚举中间值a[j]即可计算出答案,复杂度为O(len * log(len))。
总体的时间复杂度为O(block * (n/block * n/block + len * log(len))。(虽然不太懂这个复杂度,但真的很快。。。)
具体实现看代码:
#include <bits/stdc++.h>
#define fi first
#define se second
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define lowbit(x) x&-x
#define pb push_back
#define MP make_pair
#define clr(a) memset(a,0,sizeof(a))
#define _INF(a) memset(a,0x3f,sizeof(a))
#define FIN freopen("in.txt","r",stdin)
#define fuck(x) cout<<"["<<x<<"]"<<endl
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int>pii;
//head
const int MX = 1e5+5;
const double pi = acos(-1.0);
int len,mx;//开大4倍
struct Complex {
double r,i;
Complex(double r=0,double i=0):r(r),i(i) {};
Complex operator+(const Complex &rhs) {return Complex(r + rhs.r,i + rhs.i);}
Complex operator-(const Complex &rhs) {return Complex(r - rhs.r,i - rhs.i);}
Complex operator*(const Complex &rhs) {return Complex(r*rhs.r - i*rhs.i,i*rhs.r + r*rhs.i);}
} va[(1<<16)],vb[(1<<16)];
void rader(Complex F[],int len) { //len = 2^M,reverse F[i] with F[j] j为i二进制反转
int j = len >> 1;
for(int i = 1; i < len - 1; ++i) {
if(i < j) swap(F[i],F[j]); // reverse
int k = len>>1;
while(j>=k) {
j -= k;
k >>= 1;
}
if(j < k) j += k;
}
}
void FFT(Complex F[],int len,int t) {
rader(F,len);
for(int h=2; h<=len; h<<=1) {
Complex wn(cos(-t*2*pi/h),sin(-t*2*pi/h));
for(int j=0; j<len; j+=h) {
Complex E(1,0); //旋转因子
for(int k=j; k<j+h/2; ++k) {
Complex u = F[k];
Complex v = E*F[k+h/2];
F[k] = u+v;
F[k+h/2] = u-v;
E=E*wn;
}
}
}
if(t==-1) //IDFT
for(int i=0; i<len; ++i)
F[i].r/=len;
}
void Conv(Complex a[],Complex b[],int len) { //求卷积
FFT(a,len,1);
FFT(b,len,1);
for(int i=0; i<len; ++i) a[i] = a[i]*b[i];
FFT(a,len,-1);
}
int n;
int a[MX];
int Front[MX],Now[MX],Behind[MX];
void solve(){
scanf("%d",&n);
mx = 1;
for(int i = 1;i <= n;i++){
scanf("%d",&a[i]);
Behind[a[i]]++;
mx = max(mx,a[i]);
}
len = 1;
while(len <= mx) len<<=1;
len<<=1;
int block = min(n,45);
int sz = n/block;
if(n%block != 0) sz++;
ll ans = 0;
for(int pos = 1;pos <= block;pos++){
int s = sz*(pos - 1) + 1, t = min(n,sz*pos);
for(int i = s;i <= t;i++) Behind[a[i]]--;
for(int i = s;i <= t;i++){
for(int j = i + 1;j <= t;j++){
int cnt = 2*a[i] - a[j];
if(cnt>=1 && cnt <= mx){
ans += Now[cnt]; ans += Front[cnt];
}
cnt = 2*a[j] - a[i];
if(cnt>=1 && cnt <= mx){
ans += Behind[cnt];
}
}
Now[a[i]]++;
}
for(int i = 0;i < len;i++){
if(i <= mx){
va[i] = Complex(Front[i],0);
vb[i] = Complex(Behind[i],0);
}else{
va[i] = Complex(0,0);
vb[i] = Complex(0,0);
}
}
Conv(va,vb,len);
for(int i = s;i <= t;i++){
int cnt = 2*a[i];
ans += (ll)(va[cnt].r+0.5);
}
for(int i = s;i <= t;i++){
Front[a[i]]++; Now[a[i]]--;
}
}
printf("%lld\n",ans);
}
int main(){
//FIN;
solve();
return 0;
}