「bzoj3509」[CodeChef] COUNTARI - 分块+FFT | Bill Yang's Blog

「bzoj3509」[CodeChef] COUNTARI - 分块+FFT

题目大意

    给定一个长度为$N$的数组$A[]$,求有多少对$i,j,k(1\le i\lt j\lt k\le N)$满足$A[k]-A[j]=A[j]-A[i]$。


题目分析

学习到了一些FFT卡常技巧,见学习笔记

其实网上遍地的题解,我就不详细说了。
大概就是用分块平衡复杂度,使得FFT次数减小,同时块内暴力维护。


代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#include<algorithm>
#include<iostream>
#include<iomanip>
#include<cstring>
#include<cstdlib>
#include<climits>
#include<complex>
#include<vector>
#include<cstdio>
#include<cmath>
#include<queue>
using namespace std;
inline const int Get_Int() {
int num=0,bj=1;
char x=getchar();
while(!isdigit(x)) {
if(x=='-')bj=-1;
x=getchar();
}
while(isdigit(x)) {
num=num*10+x-'0';
x=getchar();
}
return num*bj;
}
#define cp complex<double>
typedef long long LL;
const int maxn=100005,maxv=30005,size_b=1800,maxb=maxn/size_b+1;
const double pi=acos(-1);
struct FastFourierTransform {
const static int maxn=262144+5;
int n,rev[maxn];
cp omega[maxn],iomega[maxn];
void init(int n) {
this->n=n;
for(int i=0; i<n; i++) {
omega[i]=cp(cos(2*pi/n*i),sin(2*pi/n*i));
iomega[i]=conj(omega[i]);
}
int k=log2(n);
for(int i=0; i<n; i++) {
int t=0;
for(int j=0; j<k; j++)if(i&(1<<j))t|=1<<(k-j-1);
rev[i]=t;
}
}
void transform(cp* a,cp* omega) {
for(int i=0; i<n; i++)if(i<rev[i])swap(a[i],a[rev[i]]);
for(int len=2; len<=n; len*=2) {
int mid=len>>1;
for(cp* p=a; p!=a+n; p+=len)
for(int i=0; i<mid; i++) {
cp t=omega[n/len*i]*p[mid+i];
p[mid+i]=p[i]-t;
p[i]+=t;
}
}
}
void dft(cp* a) {
transform(a,omega);
}
void idft(cp* a) {
transform(a,iomega);
for(int i=0; i<n; i++)a[i]=cp(a[i].real()/(n*4),a[i].imag()/n);
}
} fft;
int n,m,Max=0,BCC=0,a[maxn],Belong[maxn],L[maxb],R[maxb],Left[maxv<<1],Right[maxv<<1],Rcnt[maxv<<1];
LL Ans=0,ans[maxv<<1];
void Multiply() {
static cp c[FastFourierTransform::maxn];
fill(c,c+m,0);
for(int i=0; i<Max; i++)c[i]=cp(Left[i]+Right[i],Left[i]-Right[i]);
fft.dft(c);
for(int i=0; i<m; i++)c[i]*=c[i];
fft.idft(c);
for(int i=0; i<2*Max; i++)ans[i]=(LL)round(c[i].real());
}
void Cal(int pos) {
int val=a[pos]<<1,B=Belong[pos];
for(int i=L[B]; i<pos; i++)if(val>=a[i])Ans+=Right[val-a[i]]+Rcnt[val-a[i]];
for(int i=pos+1; i<=R[B]; i++)if(val>=a[i])Ans+=Left[val-a[i]];
if(B!=1&&B!=BCC)Ans+=ans[val];
}
int main() {
scanf("%d",&n);
for(int i=1; i<=n; i++) {
scanf("%d",&a[i]);
Max=max(Max,a[i]+1);
}
for(int i=1; i<=n; i++) {
Belong[i]=(i-1)/size_b+1;
if(!L[Belong[i]])L[Belong[i]]=i,BCC++;
R[Belong[i]]=i;
}
m=1;
while(m<(Max<<1))m<<=1;
fft.init(m);
for(int i=1; i<=n; i++)Right[a[i]]++;
for(int i=1; i<=BCC; i++) {
for(int j=L[i]; j<=R[i]; j++)Right[a[j]]--,Rcnt[a[j]]++;
if(i!=1&&i!=BCC)Multiply();
for(int j=L[i]; j<=R[i]; j++) {
Rcnt[a[j]]--;
Cal(j);
}
for(int j=L[i]; j<=R[i]; j++)Left[a[j]]++;
}
printf("%lld\n",Ans);
return 0;
}
本篇文章有用吗?有用还不快点击!
0%