#include <iostream>
#include <vector>
#include <algorithm>
int main() {
std::ios_base::sync_with_stdio(0);
long long n;
std::cin >> n;
std::vector<long long> nums(n);
for(int i = 0; i < n; ++i) {
std::cin >> nums[i];
}
std::vector<long long> pref_sums(n);
pref_sums[0] = nums[0];
for(int i = 1; i < n; ++i) {
pref_sums[i] = pref_sums[i - 1] + nums[i];
}
std::vector<long long> sums;
for(int i = 0; i < n; ++i) {
for(int j = i; j < n; ++j) {
long long sum = pref_sums[j] - (i > 0 ? pref_sums[i - 1] : 0);
sums.push_back(sum);
}
}
int N = sums.size();
std::sort(sums.begin(), sums.end());
std::vector<std::pair<long long, long long>> occurs;
long long num = sums[0];
long long occ = 1;
for(int i = 1; i < N; ++i) {
if(sums[i] == num) {
++occ;
} else {
occurs.push_back({num, occ});
num = sums[i];
occ = 1;
}
}
occurs.push_back({num, occ});
int m = occurs.size();
long long result = 0;
for(int i = 0; i < m - 2; ++i) {
int j = i + 1;
int k = m - 1;
while(j < k) {
long long sum = occurs[i].first + occurs[j].first + occurs[k].first;
if(sum == 0) {
result += occurs[i].second * occurs[j].second * occurs[k].second;
++j;
--k;
} else if(sum > 0) {
--k;
} else {
++j;
}
}
}
for(int i = 0; i < m; ++i) {
if(occurs[i].first == 0) {
long long occ = occurs[i].second;
result += occ * (occ - 1) / 2 * (occ - 2) / 3;
}
}
int i = 0;
int k = m - 1;
while(i < k) {
long long sum = 2 * occurs[i].first + occurs[k].first;
if(sum == 0) {
result += occurs[i].second * (occurs[i].second - 1) / 2 * occurs[k].second;
++i;
--k;
} else if(sum > 0) {
--k;
} else {
++i;
}
}
k = m - 1;
i = 0;
while(i < k) {
long long sum = occurs[i].first + occurs[k].first * 2;
if(sum == 0) {
result += occurs[i].second * occurs[k].second * (occurs[k].second - 1) / 2;
++i;
--k;
} else if(sum > 0) {
--k;
} else {
++i;
}
}
std::cout << result << "\n";
}
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 | #include <iostream> #include <vector> #include <algorithm> int main() { std::ios_base::sync_with_stdio(0); long long n; std::cin >> n; std::vector<long long> nums(n); for(int i = 0; i < n; ++i) { std::cin >> nums[i]; } std::vector<long long> pref_sums(n); pref_sums[0] = nums[0]; for(int i = 1; i < n; ++i) { pref_sums[i] = pref_sums[i - 1] + nums[i]; } std::vector<long long> sums; for(int i = 0; i < n; ++i) { for(int j = i; j < n; ++j) { long long sum = pref_sums[j] - (i > 0 ? pref_sums[i - 1] : 0); sums.push_back(sum); } } int N = sums.size(); std::sort(sums.begin(), sums.end()); std::vector<std::pair<long long, long long>> occurs; long long num = sums[0]; long long occ = 1; for(int i = 1; i < N; ++i) { if(sums[i] == num) { ++occ; } else { occurs.push_back({num, occ}); num = sums[i]; occ = 1; } } occurs.push_back({num, occ}); int m = occurs.size(); long long result = 0; for(int i = 0; i < m - 2; ++i) { int j = i + 1; int k = m - 1; while(j < k) { long long sum = occurs[i].first + occurs[j].first + occurs[k].first; if(sum == 0) { result += occurs[i].second * occurs[j].second * occurs[k].second; ++j; --k; } else if(sum > 0) { --k; } else { ++j; } } } for(int i = 0; i < m; ++i) { if(occurs[i].first == 0) { long long occ = occurs[i].second; result += occ * (occ - 1) / 2 * (occ - 2) / 3; } } int i = 0; int k = m - 1; while(i < k) { long long sum = 2 * occurs[i].first + occurs[k].first; if(sum == 0) { result += occurs[i].second * (occurs[i].second - 1) / 2 * occurs[k].second; ++i; --k; } else if(sum > 0) { --k; } else { ++i; } } k = m - 1; i = 0; while(i < k) { long long sum = occurs[i].first + occurs[k].first * 2; if(sum == 0) { result += occurs[i].second * occurs[k].second * (occurs[k].second - 1) / 2; ++i; --k; } else if(sum > 0) { --k; } else { ++i; } } std::cout << result << "\n"; } |
English