BZOJ3509: [CodeChef] COUNTARI
https://lydsy.com/JudgeOnline/problem.php?id=3509
分析:
- 分块,考虑\(i,j,k\)在不同的块内的方案数,这个我们枚举中间那个块两边做一次\(fft\)再枚举中间块即可。
- 其他的方案数可以用前缀桶和后缀桶求出。
代码:
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cstdlib>
#include <cmath>
using namespace std;
typedef long long ll;
typedef double f2;
#define N 200050
#define M 30000
int n,a[N],L[N],R[N],c[2][N];
const f2 pi=acos(-1);
ll ans;
struct cp {
f2 x,y;
cp() {}
cp(f2 x_,f2 y_) {x=x_, y=y_;}
cp operator + (const cp &u) const {
return cp(x+u.x, y+u.y);
}
cp operator - (const cp &u) const {
return cp(x-u.x, y-u.y);
}
cp operator * (const cp &u) const {
return cp(x*u.x-y*u.y,x*u.y+y*u.x);
}
}A[N],B[N];
void fft(cp *a,int len,int flg) {
int i,j,k,t; cp w,wn,tmp;
for(i=k=0;i<len;i++) {
if(i>k) swap(a[i],a[k]);
for(j=len>>1;(k^=j)<j;j>>=1) ;
}
for(k=2;k<=len;k<<=1) {
t=k>>1;
wn=cp(cos(2*pi*flg/k),sin(2*pi*flg/k));
for(i=0;i<len;i+=k) {
w=cp(1,0);
for(j=i;j<i+t;j++) {
tmp=a[j+t]*w;
a[j+t]=a[j]-tmp;
a[j]=a[j]+tmp;
w=w*wn;
}
}
}
if(flg==-1) for(i=0;i<len;i++) a[i].x/=len;
}
int main() {
int size=2000;
scanf("%d",&n);
int i,j,k;
for(i=1;i<=n;i++) scanf("%d",&a[i]),c[1][a[i]]++;
int blo=(n+size-1)/size;
for(i=1;i<=blo;i++) {
L[i]=R[i-1]+1, R[i]=min(n,i*size);
for(j=L[i];j<=R[i];j++) {
c[1][a[j]]--;
for(k=L[i];k<j;k++) {
if((a[j]<<1)>=a[k]) {
ans+=c[1][(a[j]<<1)-a[k]];
}
if((a[k]<<1)>=a[j]) {
ans+=c[0][(a[k]<<1)-a[j]];
}
}
}
for(j=L[i];j<=R[i];j++) c[0][a[j]]++;
}
int len=1;
while(len<=(M<<1)) len<<=1;
memset(c,0,sizeof(c));
for(i=1;i<=n;i++) c[1][a[i]]++;
for(i=L[1];i<=R[1];i++) c[0][a[i]]++,c[1][a[i]]--;
for(i=2;i<blo;i++) {
for(j=L[i];j<=R[i];j++) c[1][a[j]]--;
for(j=0;j<len;j++) A[j]=B[j]=cp(0,0);
for(j=0;j<=M;j++) A[j].x=c[0][j], B[j].x=c[1][j];
fft(A,len,1), fft(B,len,1);
for(j=0;j<len;j++) A[j]=A[j]*B[j];
fft(A,len,-1);
for(j=L[i];j<=R[i];j++) {
ans+=int(A[a[j]<<1].x+0.1);
}
for(j=L[i];j<=R[i];j++) c[0][a[j]]++;
}
printf("%lld\n",ans);
}