#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(); } |