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