#include <iostream> #include <vector> #include <unordered_map> #include <map> using namespace std; int main() { ios::sync_with_stdio(false); std::cin.tie(nullptr); short n; cin >> n; vector<short> numbers; short number; short maxAbs = 0; for (short i = 0; i < n; i++) { cin >> number; numbers.push_back(number); if (abs(number) > maxAbs) { maxAbs = abs(number); } } unordered_map<int, int> sums; vector<int> prefixes; prefixes.push_back(0); for (short i = 0; i < n; i++) { int sum = 0; for (short j = i; j < n; j++) { sum += numbers[j]; sums[sum]++; if (i == 0) { prefixes.push_back(sum); } } } int halfSize = 2 * n * maxAbs + 1; vector<vector<pair<short, int>>> beginningsAndSums(2 * halfSize); vector<vector<pair<short, int>>> endsAndSums(2 * halfSize); for (short i = 0; i < prefixes.size(); i++) { for (auto &sum : sums) { if (i < prefixes.size() - 1) { beginningsAndSums[halfSize - prefixes[i] + sum.first].push_back({i, sum.second}); } if (i > 0) { endsAndSums[halfSize + prefixes[i] + sum.first].push_back({i, sum.second}); } } } prefixes.clear(); vector<int> endsAndSumsSums(2 * halfSize); for (int i = 0; i < endsAndSums.size(); i++) { int sum = 0; for (int j = 0; j < endsAndSums[i].size(); j++) { sum += endsAndSums[i][j].second; } endsAndSumsSums[i] = sum; } long long result = 0; for (int i = 0; i < beginningsAndSums.size(); i++) { if (endsAndSums[2 * halfSize - i].size() > 0) { vector<pair<short, int>> &beginAndSum = beginningsAndSums[i]; vector<pair<short, int>> &endAndSum = endsAndSums[2 * halfSize - i]; int beginIter = 0; int endIter = 0; int sum = endsAndSumsSums[2 * halfSize - i]; while (beginIter < beginAndSum.size() && endIter < endAndSum.size()) { if (beginAndSum[beginIter].first < endAndSum[endIter].first) { result += 1LL * sum * beginAndSum[beginIter].second; beginIter++; } else { sum -= endAndSum[endIter].second; endIter++; } } } } beginningsAndSums.clear(); endsAndSums.clear(); endsAndSumsSums.clear(); for (auto &sum : sums) { if (sums.contains(-2 * sum.first)) { int doubleSum = sums[-2 * sum.first]; if (sum.first == 0) { result -= 3LL * doubleSum * (doubleSum - 1) + doubleSum; } else { result -= 3LL * sum.second * doubleSum; } } } cout << (result / 6); }
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 97 98 99 | #include <iostream> #include <vector> #include <unordered_map> #include <map> using namespace std; int main() { ios::sync_with_stdio(false); std::cin.tie(nullptr); short n; cin >> n; vector<short> numbers; short number; short maxAbs = 0; for (short i = 0; i < n; i++) { cin >> number; numbers.push_back(number); if (abs(number) > maxAbs) { maxAbs = abs(number); } } unordered_map<int, int> sums; vector<int> prefixes; prefixes.push_back(0); for (short i = 0; i < n; i++) { int sum = 0; for (short j = i; j < n; j++) { sum += numbers[j]; sums[sum]++; if (i == 0) { prefixes.push_back(sum); } } } int halfSize = 2 * n * maxAbs + 1; vector<vector<pair<short, int>>> beginningsAndSums(2 * halfSize); vector<vector<pair<short, int>>> endsAndSums(2 * halfSize); for (short i = 0; i < prefixes.size(); i++) { for (auto &sum : sums) { if (i < prefixes.size() - 1) { beginningsAndSums[halfSize - prefixes[i] + sum.first].push_back({i, sum.second}); } if (i > 0) { endsAndSums[halfSize + prefixes[i] + sum.first].push_back({i, sum.second}); } } } prefixes.clear(); vector<int> endsAndSumsSums(2 * halfSize); for (int i = 0; i < endsAndSums.size(); i++) { int sum = 0; for (int j = 0; j < endsAndSums[i].size(); j++) { sum += endsAndSums[i][j].second; } endsAndSumsSums[i] = sum; } long long result = 0; for (int i = 0; i < beginningsAndSums.size(); i++) { if (endsAndSums[2 * halfSize - i].size() > 0) { vector<pair<short, int>> &beginAndSum = beginningsAndSums[i]; vector<pair<short, int>> &endAndSum = endsAndSums[2 * halfSize - i]; int beginIter = 0; int endIter = 0; int sum = endsAndSumsSums[2 * halfSize - i]; while (beginIter < beginAndSum.size() && endIter < endAndSum.size()) { if (beginAndSum[beginIter].first < endAndSum[endIter].first) { result += 1LL * sum * beginAndSum[beginIter].second; beginIter++; } else { sum -= endAndSum[endIter].second; endIter++; } } } } beginningsAndSums.clear(); endsAndSums.clear(); endsAndSumsSums.clear(); for (auto &sum : sums) { if (sums.contains(-2 * sum.first)) { int doubleSum = sums[-2 * sum.first]; if (sum.first == 0) { result -= 3LL * doubleSum * (doubleSum - 1) + doubleSum; } else { result -= 3LL * sum.second * doubleSum; } } } cout << (result / 6); } |