//https://www.ideone.com/35iCAi #include<bits/stdc++.h> using namespace std; const int stala=510; long long tab[stala]; long long pref[stala]; vector<long long>wejscie; map<long long,long long>mapka; map<long long,long long>double_mapka; int N, M,ind=0; typedef complex<long double> cp; const long double PI = acos(-1); vector<cp> acoeff; vector<long long>result; void fft(vector<cp> &a, bool invert = false){ int n = (int)a.size(); // if n == 1 then A(w^0) = A(0) = a0 if (n == 1){ return; } //Let A_even(x) = a0+a2*x+a4*x^2+a6*x^3+... //A_odd(x) = a1+a3*x+a5*x^2+a7*x^3+... vector<cp> even, odd; //0 2 4 6 8.. for (int i = 0; i < n; i += 2){ even.push_back(a[i]); } //1 3 5 7 9.. for (int i = 1; i < n; i += 2){ odd.push_back(a[i]); } /** * divide and conquer - calculate what A_even(x^2) and A_odd(x^2) are, the sizes of each are n/2 * Put n roots of unity into A in this function, put n/2 roots of unity into the recursive one * */ fft(even, invert); fft(odd, invert); //in each iteration, w = e^(2*pi*i/n) //initially w = e^0, then multiply by e^(2*pi/n) cp w(1, 0), mult(cos(2*PI/n), sin(2*PI/n)); //in the inverted case everything is to the power of -1 if (invert) mult = {cos(-2*PI/n), sin(-2*PI/n)}; for (int i = 0; i < n/2; i++){ //A(x) = A_even(x^2)+x*A_odd(x^2), let x = w_n^i //if i < n/2, A(w_n^i) = A_even(w_(n/2)^(2*i)) + w_n^i * A_odd(w_(n/2)^(2*i)) a[i] = even[i]+w*odd[i]; //if i >= n/2, the same equation holds but i > n/2 would be out of bounds of A_even and A_odd //rearranging and substituting eq for i < n/2 yields the equation for i >= n/2 a[i+n/2] = even[i]-w*odd[i]; //in the inverted matrix, everything is divided by n at the end //since n = 2^(layers), divide the inverted matrix by 2 each time does the same thing as dividing by n at the end if (invert){ a[i] /= 2; a[i+n/2] /= 2; } //update w_n^i w *= mult; } } void prep(int n) { N=n; M=n; while (N+M > (1<<ind)){ ind++; } acoeff.resize(1<<ind); } void calculate() { fft(acoeff); vector<cp> c((1<<ind)); for (int i = 0; i < (1<<ind); i++){ c[i] = acoeff[i]*acoeff[i]; } //inverted matrix for fft - calculate C(w_n^(-i)) fft(c, true); result.resize(N+M-1); for(int i=0;i<N+M-1;i++){ result[i]=round(c[i].real()); } } int suma_pref(int x,int y) { return pref[y]-pref[x-1]; } int main() { ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0); int ile; cin>>ile; long long max_liczba=0; for(int i=1;i<=ile;i++){ cin>>tab[i]; max_liczba=max(max_liczba,abs(tab[i])); pref[i]=pref[i-1]+tab[i]; } int liczba=max_liczba*ile; int liczba2=2*liczba; prep((2*liczba)+1); wejscie.resize((2*liczba)+1); for(int i=1;i<=ile;i++){ for(int j=i;j<=ile;j++){ mapka[suma_pref(i,j)]++; double_mapka[2*suma_pref(i,j)]++; long long pom=suma_pref(i,j)+liczba; wejscie[pom]++; } } for(int i=0;i<(int)wejscie.size();i++){ acoeff[i]=wejscie[i]; } calculate(); long long wyn=0; for(int i=0;i<=liczba2*2;i++){ int x=i-liczba2; long long pom=result[i]-double_mapka[x]; result[i]=result[i]-(pom/2); wyn+=(mapka[-x]*result[i]); } for(int i=1;i<=ile;i++){ for(int j=i;j<=ile;j++){ long long pom=suma_pref(i,j)*2; wyn-=(2*mapka[-pom]); if(pom==0){ wyn++; } } } cout<<wyn/3; 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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | //https://www.ideone.com/35iCAi #include<bits/stdc++.h> using namespace std; const int stala=510; long long tab[stala]; long long pref[stala]; vector<long long>wejscie; map<long long,long long>mapka; map<long long,long long>double_mapka; int N, M,ind=0; typedef complex<long double> cp; const long double PI = acos(-1); vector<cp> acoeff; vector<long long>result; void fft(vector<cp> &a, bool invert = false){ int n = (int)a.size(); // if n == 1 then A(w^0) = A(0) = a0 if (n == 1){ return; } //Let A_even(x) = a0+a2*x+a4*x^2+a6*x^3+... //A_odd(x) = a1+a3*x+a5*x^2+a7*x^3+... vector<cp> even, odd; //0 2 4 6 8.. for (int i = 0; i < n; i += 2){ even.push_back(a[i]); } //1 3 5 7 9.. for (int i = 1; i < n; i += 2){ odd.push_back(a[i]); } /** * divide and conquer - calculate what A_even(x^2) and A_odd(x^2) are, the sizes of each are n/2 * Put n roots of unity into A in this function, put n/2 roots of unity into the recursive one * */ fft(even, invert); fft(odd, invert); //in each iteration, w = e^(2*pi*i/n) //initially w = e^0, then multiply by e^(2*pi/n) cp w(1, 0), mult(cos(2*PI/n), sin(2*PI/n)); //in the inverted case everything is to the power of -1 if (invert) mult = {cos(-2*PI/n), sin(-2*PI/n)}; for (int i = 0; i < n/2; i++){ //A(x) = A_even(x^2)+x*A_odd(x^2), let x = w_n^i //if i < n/2, A(w_n^i) = A_even(w_(n/2)^(2*i)) + w_n^i * A_odd(w_(n/2)^(2*i)) a[i] = even[i]+w*odd[i]; //if i >= n/2, the same equation holds but i > n/2 would be out of bounds of A_even and A_odd //rearranging and substituting eq for i < n/2 yields the equation for i >= n/2 a[i+n/2] = even[i]-w*odd[i]; //in the inverted matrix, everything is divided by n at the end //since n = 2^(layers), divide the inverted matrix by 2 each time does the same thing as dividing by n at the end if (invert){ a[i] /= 2; a[i+n/2] /= 2; } //update w_n^i w *= mult; } } void prep(int n) { N=n; M=n; while (N+M > (1<<ind)){ ind++; } acoeff.resize(1<<ind); } void calculate() { fft(acoeff); vector<cp> c((1<<ind)); for (int i = 0; i < (1<<ind); i++){ c[i] = acoeff[i]*acoeff[i]; } //inverted matrix for fft - calculate C(w_n^(-i)) fft(c, true); result.resize(N+M-1); for(int i=0;i<N+M-1;i++){ result[i]=round(c[i].real()); } } int suma_pref(int x,int y) { return pref[y]-pref[x-1]; } int main() { ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0); int ile; cin>>ile; long long max_liczba=0; for(int i=1;i<=ile;i++){ cin>>tab[i]; max_liczba=max(max_liczba,abs(tab[i])); pref[i]=pref[i-1]+tab[i]; } int liczba=max_liczba*ile; int liczba2=2*liczba; prep((2*liczba)+1); wejscie.resize((2*liczba)+1); for(int i=1;i<=ile;i++){ for(int j=i;j<=ile;j++){ mapka[suma_pref(i,j)]++; double_mapka[2*suma_pref(i,j)]++; long long pom=suma_pref(i,j)+liczba; wejscie[pom]++; } } for(int i=0;i<(int)wejscie.size();i++){ acoeff[i]=wejscie[i]; } calculate(); long long wyn=0; for(int i=0;i<=liczba2*2;i++){ int x=i-liczba2; long long pom=result[i]-double_mapka[x]; result[i]=result[i]-(pom/2); wyn+=(mapka[-x]*result[i]); } for(int i=1;i<=ile;i++){ for(int j=i;j<=ile;j++){ long long pom=suma_pref(i,j)*2; wyn-=(2*mapka[-pom]); if(pom==0){ wyn++; } } } cout<<wyn/3; return 0; } |