#include <bits/stdc++.h> using namespace std; typedef double LD; typedef complex<LD> CD; int N; const LD PI=acos(-1); void fft(vector<CD> & a, bool invert){ int n=a.size(); for(int i=1, j=0; i<n; i++){ int bit=n>>1; for(; j&bit; bit>>=1) j^=bit; j^=bit; if(i<j) swap(a[i], a[j]); } for(int len=2; len<=n; len<<=1){ LD ang=2*PI/len*(invert ? -1 : 1); CD wlen(cos(ang), sin(ang)); for(int i=0; i<n; i+=len){ CD w(1); for(int j=0; j<len/2; j++){ CD u=a[i+j], v=a[i+j+len/2]*w; a[i+j]=u+v; a[i+j+len/2]=u-v; w*=wlen; } } } if(invert) { for(CD&x : a) x/=n; } } int multiply(vector<int> const& a, vector<int>& num) { vector<CD> fa(a.begin(), a.end()); int n=a.size()*2; fa.resize(n); fft(fa, false); for(int i=0; i<n; i++) fa[i]*=fa[i]; fft(fa, true); int ans=0; for(int i : num) ans+=round(fa[n/2-i].real()); for(int i : num) ans-=a[n/4+(-2*i)]*3; ans+=a[n/4]*2; return ans/6; } int THREESUM(vector<int>& nums) { int n=1; while(n<N) n<<=1; vector<int> present(n*4, 0); for(int num : nums) present[num+2*n]++; return multiply(present, nums); } int main() { ios_base::sync_with_stdio(0); cin.tie(0); int n; cin>>n; int a; vector<int> C1; for(int i=0; i<n; i++){ cin>>a; C1.push_back(a); } vector<int> C2; int minim=99999999; int maxim=-99999999; for(int i=0; i<n; i++){ int sum=0; for(int j=i; j<n; j++){ sum+=C1[j]; C2.push_back(sum); maxim=max(maxim, sum); minim=min(minim, sum); } } /*int p=(maxim+minim)/2; for(int i=0; i<C2.size(); i++){ C2[i]-=p; }*/ N=max(abs(maxim), abs(minim))+2; cout<<THREESUM(C2); return 0; }
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 | #include <bits/stdc++.h> using namespace std; typedef double LD; typedef complex<LD> CD; int N; const LD PI=acos(-1); void fft(vector<CD> & a, bool invert){ int n=a.size(); for(int i=1, j=0; i<n; i++){ int bit=n>>1; for(; j&bit; bit>>=1) j^=bit; j^=bit; if(i<j) swap(a[i], a[j]); } for(int len=2; len<=n; len<<=1){ LD ang=2*PI/len*(invert ? -1 : 1); CD wlen(cos(ang), sin(ang)); for(int i=0; i<n; i+=len){ CD w(1); for(int j=0; j<len/2; j++){ CD u=a[i+j], v=a[i+j+len/2]*w; a[i+j]=u+v; a[i+j+len/2]=u-v; w*=wlen; } } } if(invert) { for(CD&x : a) x/=n; } } int multiply(vector<int> const& a, vector<int>& num) { vector<CD> fa(a.begin(), a.end()); int n=a.size()*2; fa.resize(n); fft(fa, false); for(int i=0; i<n; i++) fa[i]*=fa[i]; fft(fa, true); int ans=0; for(int i : num) ans+=round(fa[n/2-i].real()); for(int i : num) ans-=a[n/4+(-2*i)]*3; ans+=a[n/4]*2; return ans/6; } int THREESUM(vector<int>& nums) { int n=1; while(n<N) n<<=1; vector<int> present(n*4, 0); for(int num : nums) present[num+2*n]++; return multiply(present, nums); } int main() { ios_base::sync_with_stdio(0); cin.tie(0); int n; cin>>n; int a; vector<int> C1; for(int i=0; i<n; i++){ cin>>a; C1.push_back(a); } vector<int> C2; int minim=99999999; int maxim=-99999999; for(int i=0; i<n; i++){ int sum=0; for(int j=i; j<n; j++){ sum+=C1[j]; C2.push_back(sum); maxim=max(maxim, sum); minim=min(minim, sum); } } /*int p=(maxim+minim)/2; for(int i=0; i<C2.size(); i++){ C2[i]-=p; }*/ N=max(abs(maxim), abs(minim))+2; cout<<THREESUM(C2); return 0; } |