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
#include <cstdio>
#include <cstdlib>
#include <cstdint>

#include <vector>
#include <algorithm>

int main() {
    int n;
    scanf("%d", &n);

    std::vector<int> nums, dnums;
    nums.reserve(n);
    dnums.reserve((n * (n + 1)) / 2);

    for (int i = 0; i < n; i++) {
        int x;
        scanf("%d", &x);
        nums.push_back(x);
        int sum = 0;
        for (int j = i; j >= 0; j--) {
            sum += nums[j];
            // printf("sum %i, %i: %i\n", j, i, sum);
            dnums.push_back(sum);
        }
    }

    std::sort(dnums.begin(), dnums.end());

    std::vector<std::pair<int, int>> gnums;
    int count_so_far = 0;
    int prev = dnums[0];
    for (int x : dnums) {
        if (x == prev) {
            count_so_far++;
        } else {
            gnums.push_back({prev, count_so_far});
            prev = x;
            count_so_far = 1;
        }
    }
    gnums.push_back({prev, count_so_far});

    // for (auto [s, c] : gnums) {
    //     printf(" sum: %d, count: %lld\n", s, c);
    // }

    int total_count = 0;
    for (int i = 0; i < gnums.size(); i++) {
        const auto& [target_sum, sum_count] = gnums[i];
        // printf(" targetting %d (occurrences: %lld)\n", target_sum, sum_count);
        int a = i;
        int b = gnums.size() - 1;
        int local_count = 0;
        while (a <= b) {
            const int sum = gnums[a].first + gnums[b].first;
            // printf("  considering %d, %d (%lld)\n", gnums[a].first, gnums[b].first, sum + target_sum);
            if (sum == -target_sum) {
                int increment;
                if (i != a && i != b && a != b) {
                    // All different
                    increment = sum_count * gnums[a].second * gnums[b].second;
                } else if (i == a && i == b) {
                    // All equal
                    increment = (sum_count * (sum_count - 1) * (sum_count - 2)) / 6;
                } else if (a == b) {
                    // Only a == b
                    increment = sum_count * (gnums[a].second * (gnums[a].second - 1)) / 2;
                } else if (i == a) {
                    // Only i == a
                    increment = gnums[b].second * (sum_count * (sum_count - 1)) / 2;
                } else {
                    // Only i == b
                    increment = gnums[a].second * (sum_count * (sum_count - 1)) / 2;
                }
                // printf("   adding %lld\n", increment);
                local_count += increment;
            }
            if (sum > -target_sum) {
                b--;
            } else {
                a++;
            }
        }
        total_count += local_count;
    }

    printf("%d\n", total_count);

    return 0;
}