// clang-format off #include<bits/stdc++.h> using namespace std; using LL=long long; #define FOR(i,l,r) for(auto i=(l);i<=(r);++i) #define REP(i,n) FOR(i,0,(n)-1) #define ssize(x) int(x.size()) template<class A,class B>auto&operator<<(ostream&o,pair<A,B>p){return o<<"("<<p.first<<", "<<p.second<<")";} template<class T>auto operator<<(ostream&o,T x)->decltype(x.end(),o){o<<"{";int i=0;for(auto e:x)o<<(", ")+2*!i++<<e;return o<<"}";} #ifdef DEBUG #define debug(x...) cerr<<"["#x"]: ",[](auto...$){((cerr<<$<<"; "),...);}(x),cerr<<'\n' #else #define debug(...) {} #endif // clang-format on using cd = complex<double>; 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) { 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; } } void gen_conv(const unordered_map<int, int>& Bcnt, vector<LL>& out) { vector<cd> conv(ssize(out)); for (const auto& x : Bcnt) conv[x.first] = x.second; fft(conv, false); for (auto& x : conv) x *= x; fft(conv, true); REP (i, ssize(out)) out[i] = round(conv[i].real()); } int main() { cin.tie(0)->sync_with_stdio(0); int n, nB; cin >> n; vector<int> A(n), B; for (auto& x : A) cin >> x; // Generate B sequence int Bmin = numeric_limits<int>::max(), Bmax = numeric_limits<int>::min(); B.reserve(nB = (n * (n + 1)) / 2); REP (start, n) { int sum = 0; FOR (i, start, n - 1) { sum += A[i]; B.emplace_back(sum); Bmin = min(Bmin, sum); Bmax = max(Bmax, sum); } } debug(B, Bmin, Bmax); int base = max(0, -Bmin) + 1, conv_size = 1; while (conv_size <= base + Bmax) conv_size <<= 1; conv_size <<= 1; // Find the count of individual values unordered_map<int, int> single_cnt; for (auto& x : B) { x += base; single_cnt[x] += 1; } debug(B, base, conv_size); // Find the number of pairs that give given sum vector<LL> cnt(conv_size); gen_conv(single_cnt, cnt); for (const auto& x : single_cnt) if (x.second >= 2) cnt[x.first << 1] -= x.second; debug(cnt); // Calculate the score LL score = 0; const int my_zero = base * 3; for (const auto& x : B) { const int& sum = my_zero - x; if (sum >= conv_size || sum <= 0) continue; score += cnt[sum] / 2; // Subtract the case, when cnt[my_zero -x] contains x const auto& it = single_cnt.find(sum - x); if (it == single_cnt.end()) continue; score -= it->second; if (x == base) score += 1; } cout << score / 3 << "\n"; #ifdef LOCAL system("grep VmPeak /proc/$PPID/status >&2"); #endif 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 | // clang-format off #include<bits/stdc++.h> using namespace std; using LL=long long; #define FOR(i,l,r) for(auto i=(l);i<=(r);++i) #define REP(i,n) FOR(i,0,(n)-1) #define ssize(x) int(x.size()) template<class A,class B>auto&operator<<(ostream&o,pair<A,B>p){return o<<"("<<p.first<<", "<<p.second<<")";} template<class T>auto operator<<(ostream&o,T x)->decltype(x.end(),o){o<<"{";int i=0;for(auto e:x)o<<(", ")+2*!i++<<e;return o<<"}";} #ifdef DEBUG #define debug(x...) cerr<<"["#x"]: ",[](auto...$){((cerr<<$<<"; "),...);}(x),cerr<<'\n' #else #define debug(...) {} #endif // clang-format on using cd = complex<double>; 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) { 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; } } void gen_conv(const unordered_map<int, int>& Bcnt, vector<LL>& out) { vector<cd> conv(ssize(out)); for (const auto& x : Bcnt) conv[x.first] = x.second; fft(conv, false); for (auto& x : conv) x *= x; fft(conv, true); REP (i, ssize(out)) out[i] = round(conv[i].real()); } int main() { cin.tie(0)->sync_with_stdio(0); int n, nB; cin >> n; vector<int> A(n), B; for (auto& x : A) cin >> x; // Generate B sequence int Bmin = numeric_limits<int>::max(), Bmax = numeric_limits<int>::min(); B.reserve(nB = (n * (n + 1)) / 2); REP (start, n) { int sum = 0; FOR (i, start, n - 1) { sum += A[i]; B.emplace_back(sum); Bmin = min(Bmin, sum); Bmax = max(Bmax, sum); } } debug(B, Bmin, Bmax); int base = max(0, -Bmin) + 1, conv_size = 1; while (conv_size <= base + Bmax) conv_size <<= 1; conv_size <<= 1; // Find the count of individual values unordered_map<int, int> single_cnt; for (auto& x : B) { x += base; single_cnt[x] += 1; } debug(B, base, conv_size); // Find the number of pairs that give given sum vector<LL> cnt(conv_size); gen_conv(single_cnt, cnt); for (const auto& x : single_cnt) if (x.second >= 2) cnt[x.first << 1] -= x.second; debug(cnt); // Calculate the score LL score = 0; const int my_zero = base * 3; for (const auto& x : B) { const int& sum = my_zero - x; if (sum >= conv_size || sum <= 0) continue; score += cnt[sum] / 2; // Subtract the case, when cnt[my_zero -x] contains x const auto& it = single_cnt.find(sum - x); if (it == single_cnt.end()) continue; score -= it->second; if (x == base) score += 1; } cout << score / 3 << "\n"; #ifdef LOCAL system("grep VmPeak /proc/$PPID/status >&2"); #endif return 0; } |