#include <bits/stdc++.h> using namespace std; int main() { cin.tie(0)->sync_with_stdio(false); int n; cin >> n; vector<int> uc(n); for (int &x : uc) cin >> x; vector<int> buc; buc.reserve(n * (n + 1) / 2); for (int i = 0; i < n; i++) { for (int ts = 0, j = i; j < n; j++) buc.push_back(ts += uc[j]); } sort(buc.begin(), buc.end()); // copy(buc.begin(), buc.end(), ostream_iterator<int>(cout, " ")); cout << endl; int m = buc.size(); long long cnt = 0; vector<long long> ubuc, ubcnt; unordered_map<int, long long> ucm; ubuc.reserve(m); long long p = buc[0] - 1; // compress for (int i = 0; i < m; i++) { if (buc[i] != p) ubuc.push_back(p = buc[i]); ucm[p] += 1; } int k = ubuc.size(); ubcnt.resize(k); for (int i = 0; i < k; i++) ubcnt[i] = ucm[ubuc[i]]; // 1sum = many zeroes? if (ucm[0] >= 3) { cnt = ucm[0] * (ucm[0] - 1) * (ucm[0] - 2) / 6; } // 2sum for (int i = 0; i < k; i++) { if (ubuc[i] == 0) continue; long long cl = ubcnt[i]; if (cl < 2) continue; long long sv = -2 * ubuc[i]; long long cr = ucm[sv]; cnt += cl * (cl - 1) * cr / 2; } // 3sum for (int i = 0; i < k - 2; i++) { int l = i + 1, r = k - 1; if (ubuc[i] + ubuc[l] > 0) continue; while (l < r) { long long cm = ubcnt[i]; long long sum3 = ubuc[i] + ubuc[l] + ubuc[r]; if (sum3 == 0) { long long cl = ubcnt[l]; long long cr = ubcnt[r]; // cout << ubuc[i] << "*" << cm << "," << ubuc[l] << "*" << cl << "," << ubuc[r] << "*" << cr << endl; cnt += cl * cr * cm; l += 1; r -= 1; } else if (sum3 < 0) { l += 1; } else { r -= 1; } } } cout << cnt << endl; 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 | #include <bits/stdc++.h> using namespace std; int main() { cin.tie(0)->sync_with_stdio(false); int n; cin >> n; vector<int> uc(n); for (int &x : uc) cin >> x; vector<int> buc; buc.reserve(n * (n + 1) / 2); for (int i = 0; i < n; i++) { for (int ts = 0, j = i; j < n; j++) buc.push_back(ts += uc[j]); } sort(buc.begin(), buc.end()); // copy(buc.begin(), buc.end(), ostream_iterator<int>(cout, " ")); cout << endl; int m = buc.size(); long long cnt = 0; vector<long long> ubuc, ubcnt; unordered_map<int, long long> ucm; ubuc.reserve(m); long long p = buc[0] - 1; // compress for (int i = 0; i < m; i++) { if (buc[i] != p) ubuc.push_back(p = buc[i]); ucm[p] += 1; } int k = ubuc.size(); ubcnt.resize(k); for (int i = 0; i < k; i++) ubcnt[i] = ucm[ubuc[i]]; // 1sum = many zeroes? if (ucm[0] >= 3) { cnt = ucm[0] * (ucm[0] - 1) * (ucm[0] - 2) / 6; } // 2sum for (int i = 0; i < k; i++) { if (ubuc[i] == 0) continue; long long cl = ubcnt[i]; if (cl < 2) continue; long long sv = -2 * ubuc[i]; long long cr = ucm[sv]; cnt += cl * (cl - 1) * cr / 2; } // 3sum for (int i = 0; i < k - 2; i++) { int l = i + 1, r = k - 1; if (ubuc[i] + ubuc[l] > 0) continue; while (l < r) { long long cm = ubcnt[i]; long long sum3 = ubuc[i] + ubuc[l] + ubuc[r]; if (sum3 == 0) { long long cl = ubcnt[l]; long long cr = ubcnt[r]; // cout << ubuc[i] << "*" << cm << "," << ubuc[l] << "*" << cl << "," << ubuc[r] << "*" << cr << endl; cnt += cl * cr * cm; l += 1; r -= 1; } else if (sum3 < 0) { l += 1; } else { r -= 1; } } } cout << cnt << endl; return 0; } |