#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; } |
English