#include <algorithm> #include <numeric> #include <iostream> #include <unordered_map> #include <unordered_set> #include <vector> using namespace std; struct Triple { short a_b, c; }; int main() { int n; cin >> n; vector<int> nums, prefixes; nums.resize(n + 1); for (int i = 1; i <= n; ++i) { cin >> nums[i]; } prefixes.resize(n + 1); partial_sum(begin(nums), end(nums), begin(prefixes)); unordered_map<int, vector<Triple>> starts, ends; for (short k = 0; k < n; ++k) { for (short i = 0; i < n; ++i) { for (short j = i + 1; j <= n; ++j) { int sKey = -prefixes[i] + prefixes[j] - prefixes[k]; auto &s = starts[sKey]; if (size(s) == 0 || rbegin(s)->c != k) { s.push_back({0, k}); } ++s.back().a_b; int eKey = -prefixes[i] + prefixes[j] + prefixes[k + 1]; auto &e = ends[eKey]; if (size(e) == 0 || rbegin(e)->c != k + 1) { e.push_back({0, short(k + 1)}); } ++e.back().a_b; } } } long long res{}; for (auto &start : starts) { auto &s = start.second; if (!ends.contains(-start.first)) continue; auto &e = ends[-start.first]; int i{}, j{}; long long jTotal = accumulate(begin(e), end(e), 0, [](long long acc, const Triple &v) { return acc + v.a_b;}); long long jAcc{}; while (i < size(s) && j < size(e)) { while (j < size(e) && s[i].c >= e[j].c) { jAcc += e[j].a_b; ++j; } if (j >= size(e)) break; res += s[i].a_b * (jTotal - jAcc); ++i; } } unordered_map<int, int> sizes; for (int i = 0; i < n; ++i) { for (int j = i + 1; j <= n; ++j) { ++sizes[prefixes[j] - prefixes[i]]; } } for (auto &s : sizes) { if (s.first != 0) { if (sizes.contains(-2 * s.first)) { res -= 3LL * s.second * sizes[-2 * s.first]; } } else { res -= s.second + 3LL * s.second * (s.second - 1); } } cout << res / 6 << endl; }
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 | #include <algorithm> #include <numeric> #include <iostream> #include <unordered_map> #include <unordered_set> #include <vector> using namespace std; struct Triple { short a_b, c; }; int main() { int n; cin >> n; vector<int> nums, prefixes; nums.resize(n + 1); for (int i = 1; i <= n; ++i) { cin >> nums[i]; } prefixes.resize(n + 1); partial_sum(begin(nums), end(nums), begin(prefixes)); unordered_map<int, vector<Triple>> starts, ends; for (short k = 0; k < n; ++k) { for (short i = 0; i < n; ++i) { for (short j = i + 1; j <= n; ++j) { int sKey = -prefixes[i] + prefixes[j] - prefixes[k]; auto &s = starts[sKey]; if (size(s) == 0 || rbegin(s)->c != k) { s.push_back({0, k}); } ++s.back().a_b; int eKey = -prefixes[i] + prefixes[j] + prefixes[k + 1]; auto &e = ends[eKey]; if (size(e) == 0 || rbegin(e)->c != k + 1) { e.push_back({0, short(k + 1)}); } ++e.back().a_b; } } } long long res{}; for (auto &start : starts) { auto &s = start.second; if (!ends.contains(-start.first)) continue; auto &e = ends[-start.first]; int i{}, j{}; long long jTotal = accumulate(begin(e), end(e), 0, [](long long acc, const Triple &v) { return acc + v.a_b;}); long long jAcc{}; while (i < size(s) && j < size(e)) { while (j < size(e) && s[i].c >= e[j].c) { jAcc += e[j].a_b; ++j; } if (j >= size(e)) break; res += s[i].a_b * (jTotal - jAcc); ++i; } } unordered_map<int, int> sizes; for (int i = 0; i < n; ++i) { for (int j = i + 1; j <= n; ++j) { ++sizes[prefixes[j] - prefixes[i]]; } } for (auto &s : sizes) { if (s.first != 0) { if (sizes.contains(-2 * s.first)) { res -= 3LL * s.second * sizes[-2 * s.first]; } } else { res -= s.second + 3LL * s.second * (s.second - 1); } } cout << res / 6 << endl; } |