#include <bits/stdc++.h>
using namespace std;
typedef complex<long double> cd;
const double 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) {
long double 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 main(){
int n;
cin >> n;
vector <int> A(n);
for(auto &a: A) cin >> a;
vector <int> S;
int min_val = A[0], max_val = A[0];
map <int, int> M;
for(int i = 0; i<n; ++i){
int sum = 0;
for(int j = i; j<n; ++j){
sum += A[j];
S.push_back(sum);
M[sum]++;
min_val = min(min_val, sum);
max_val = max(max_val, sum);
}
}
long res = 0;
res -= M[0]*M[0]*M[0] - M[0]*(M[0]-1)*(M[0]-2);
for(auto m: M){
if(m.first != 0){
int y = -2*m.first;
res -= 3*m.second*m.second*M[y];
res += 3*m.second*(m.second-1)*M[y];
//cout << "RES: " << res << " " << m.first << "\n";
}
}
int N = max_val - min_val + 1;
if(min_val > 0 || max_val < 0){
cout << 0 << "\n";
return 0;
}
for(int i = 0; i<=25; ++i){
if(N <= (1<<i)){
N = (1<<i);
break;
}
}
vector <cd> P(N);
for(auto s: S){
int act_s = s - min_val;
P[act_s] = P[act_s] + (cd)1;
}
//for(auto p: P) cout << p << "\n";
for(int i = 0; i<3*N; ++i){
P.push_back(0);
}
fft(P, 0);
vector <cd> Q(4*N);
for(int i = 0; i<4*N; ++i){
Q[i] = P[i]*P[i]*P[i];
}
fft(Q, 1);
//for(auto q: Q) cout << q << " ";
//cout << "\n";
long long v = ceil(Q[-3*min_val].real());
cout << (res + v)/6 << "\n";
}
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 <bits/stdc++.h> using namespace std; typedef complex<long double> cd; const double 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) { long double 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 main(){ int n; cin >> n; vector <int> A(n); for(auto &a: A) cin >> a; vector <int> S; int min_val = A[0], max_val = A[0]; map <int, int> M; for(int i = 0; i<n; ++i){ int sum = 0; for(int j = i; j<n; ++j){ sum += A[j]; S.push_back(sum); M[sum]++; min_val = min(min_val, sum); max_val = max(max_val, sum); } } long res = 0; res -= M[0]*M[0]*M[0] - M[0]*(M[0]-1)*(M[0]-2); for(auto m: M){ if(m.first != 0){ int y = -2*m.first; res -= 3*m.second*m.second*M[y]; res += 3*m.second*(m.second-1)*M[y]; //cout << "RES: " << res << " " << m.first << "\n"; } } int N = max_val - min_val + 1; if(min_val > 0 || max_val < 0){ cout << 0 << "\n"; return 0; } for(int i = 0; i<=25; ++i){ if(N <= (1<<i)){ N = (1<<i); break; } } vector <cd> P(N); for(auto s: S){ int act_s = s - min_val; P[act_s] = P[act_s] + (cd)1; } //for(auto p: P) cout << p << "\n"; for(int i = 0; i<3*N; ++i){ P.push_back(0); } fft(P, 0); vector <cd> Q(4*N); for(int i = 0; i<4*N; ++i){ Q[i] = P[i]*P[i]*P[i]; } fft(Q, 1); //for(auto q: Q) cout << q << " "; //cout << "\n"; long long v = ceil(Q[-3*min_val].real()); cout << (res + v)/6 << "\n"; } |
English