#include <bits/stdc++.h>
using namespace std;
#define JOIN_(X, Y) X##Y
#define JOIN(X, Y) JOIN_(X, Y)
#define TMP JOIN(tmp, __LINE__)
#define PB push_back
#define SZ(x) int((x).size())
#define REP(i, n) for (int i = 0, TMP = (n); i < TMP; ++i)
#define FOR(i, a, b) for (int i = (a), TMP = (b); i <= TMP; ++i)
#define FORD(i, a, b) for (int i = (a), TMP = (b); i >= TMP; --i)
#ifdef DEBUG
#define DEB(x) (cerr << x)
#else
#define DEB(x)
#endif
typedef long long ll;
typedef vector<int> vi;
typedef pair<int, int> pii;
const ll INF = static_cast<ll>(1e18) + 9;
void inline one() {
int n;
cin >> n;
vi prefix_sum(1, 0);
REP(i, n) {
int x;
cin >> x;
int sum = prefix_sum.back() + x;
prefix_sum.PB(sum);
}
vi t;
REP(i, n) {
FOR(j, i + 1, n) { t.PB(prefix_sum[j] - prefix_sum[i]); }
}
unordered_map<int, ll> cnt;
DEB("t:\n");
for (int x : t) {
DEB(x << " ");
++cnt[x];
}
DEB("\n");
DEB("cnt:\n");
ll result = 0;
vector<pair<int, ll>> pairs;
for (const auto &[k, v] : cnt) {
// DEB(k << ": " << v << "\n");
pairs.emplace_back(k, v);
}
DEB("SZ(t)=" << SZ(t) << " SZ(cnt)=" << SZ(cnt) << " SZ(pairs)=" << SZ(pairs) << "\n");
sort(pairs.begin(), pairs.end());
FOR(i, 0, SZ(pairs) - 1) {
auto [a, a_cnt] = pairs[i];
// DEB("a=" << a << "\n");
if (a > 0) {
break;
}
// int start = i;
// int end = SZ(pairs) - 1;
const auto &ub = upper_bound(pairs.begin() + i, pairs.end(), pair<int, ll>(-2 * a, INF));
int end = distance(pairs.begin(), ub) - 1;
const auto &lb = lower_bound(pairs.begin() + i, pairs.end(),
pair<int, ll>(-(a + pairs[end].first), -INF));
int start = distance(pairs.begin(), lb);
while (start <= end) {
auto [b, b_cnt] = pairs[start];
auto [c, c_cnt] = pairs[end];
int sum = a + b + c;
if (sum == 0) {
if (i == start and start == end) {
if (a_cnt >= 3) {
result += a_cnt * (b_cnt - 1) * (c_cnt - 2) / 6;
}
} else if (i == start) {
if (a_cnt >= 2) {
result += (a_cnt * (b_cnt - 1)) / 2 * c_cnt;
}
} else if (start == end) {
if (b_cnt >= 2) {
result += a_cnt * (b_cnt * (c_cnt - 1)) / 2;
}
} else {
result += a_cnt * b_cnt * c_cnt;
}
++start;
--end;
} else if (sum > 0) {
--end;
} else {
++start;
}
}
}
cout << result << "\n";
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
// int z; cin >> z; while(z--)
one();
}
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 | #include <bits/stdc++.h> using namespace std; #define JOIN_(X, Y) X##Y #define JOIN(X, Y) JOIN_(X, Y) #define TMP JOIN(tmp, __LINE__) #define PB push_back #define SZ(x) int((x).size()) #define REP(i, n) for (int i = 0, TMP = (n); i < TMP; ++i) #define FOR(i, a, b) for (int i = (a), TMP = (b); i <= TMP; ++i) #define FORD(i, a, b) for (int i = (a), TMP = (b); i >= TMP; --i) #ifdef DEBUG #define DEB(x) (cerr << x) #else #define DEB(x) #endif typedef long long ll; typedef vector<int> vi; typedef pair<int, int> pii; const ll INF = static_cast<ll>(1e18) + 9; void inline one() { int n; cin >> n; vi prefix_sum(1, 0); REP(i, n) { int x; cin >> x; int sum = prefix_sum.back() + x; prefix_sum.PB(sum); } vi t; REP(i, n) { FOR(j, i + 1, n) { t.PB(prefix_sum[j] - prefix_sum[i]); } } unordered_map<int, ll> cnt; DEB("t:\n"); for (int x : t) { DEB(x << " "); ++cnt[x]; } DEB("\n"); DEB("cnt:\n"); ll result = 0; vector<pair<int, ll>> pairs; for (const auto &[k, v] : cnt) { // DEB(k << ": " << v << "\n"); pairs.emplace_back(k, v); } DEB("SZ(t)=" << SZ(t) << " SZ(cnt)=" << SZ(cnt) << " SZ(pairs)=" << SZ(pairs) << "\n"); sort(pairs.begin(), pairs.end()); FOR(i, 0, SZ(pairs) - 1) { auto [a, a_cnt] = pairs[i]; // DEB("a=" << a << "\n"); if (a > 0) { break; } // int start = i; // int end = SZ(pairs) - 1; const auto &ub = upper_bound(pairs.begin() + i, pairs.end(), pair<int, ll>(-2 * a, INF)); int end = distance(pairs.begin(), ub) - 1; const auto &lb = lower_bound(pairs.begin() + i, pairs.end(), pair<int, ll>(-(a + pairs[end].first), -INF)); int start = distance(pairs.begin(), lb); while (start <= end) { auto [b, b_cnt] = pairs[start]; auto [c, c_cnt] = pairs[end]; int sum = a + b + c; if (sum == 0) { if (i == start and start == end) { if (a_cnt >= 3) { result += a_cnt * (b_cnt - 1) * (c_cnt - 2) / 6; } } else if (i == start) { if (a_cnt >= 2) { result += (a_cnt * (b_cnt - 1)) / 2 * c_cnt; } } else if (start == end) { if (b_cnt >= 2) { result += a_cnt * (b_cnt * (c_cnt - 1)) / 2; } } else { result += a_cnt * b_cnt * c_cnt; } ++start; --end; } else if (sum > 0) { --end; } else { ++start; } } } cout << result << "\n"; } int main() { ios::sync_with_stdio(false); cin.tie(0); // int z; cin >> z; while(z--) one(); } |
English