from collections import defaultdict from math import comb def countDoubles(counter): count = 0 for ele in counter.keys(): if ele != 0 and ele in counter and counter[ele] >= 2 and ele*-2 in counter and counter[ele*-2] > 0: count += comb(counter[ele], 2) * counter[ele*-2] return count n = int(input()) sequence = [int(x) for x in input().split()] positive_counts = defaultdict(int) negative_counts = defaultdict(int) counts = defaultdict(int) for i in range(n): acc = 0 for j in range(n-i): acc += sequence[i+j] counts[acc] += 1 if acc >= 0: positive_counts[acc] += 1 elif acc < 0: negative_counts[acc] += 1 count_triplets = 0 positive_keys = list(positive_counts.keys()) negative_keys = list(negative_counts.keys()) for i in range(len(positive_keys)): for j in range(i+1, len(positive_keys)): target = -(positive_keys[i] + positive_keys[j]) if target in negative_counts: count_triplets += positive_counts[positive_keys[i]] * \ positive_counts[positive_keys[j]] * negative_counts[target] for i in range(len(negative_keys)): for j in range(i+1, len(negative_keys)): target = -(negative_keys[i] + negative_keys[j]) if target in positive_counts: count_triplets += negative_counts[negative_keys[i]] * \ negative_counts[negative_keys[j]] * positive_counts[target] print(count_triplets + countDoubles(counts) + comb(counts[0], 3))
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 | from collections import defaultdict from math import comb def countDoubles(counter): count = 0 for ele in counter.keys(): if ele != 0 and ele in counter and counter[ele] >= 2 and ele*-2 in counter and counter[ele*-2] > 0: count += comb(counter[ele], 2) * counter[ele*-2] return count n = int(input()) sequence = [int(x) for x in input().split()] positive_counts = defaultdict(int) negative_counts = defaultdict(int) counts = defaultdict(int) for i in range(n): acc = 0 for j in range(n-i): acc += sequence[i+j] counts[acc] += 1 if acc >= 0: positive_counts[acc] += 1 elif acc < 0: negative_counts[acc] += 1 count_triplets = 0 positive_keys = list(positive_counts.keys()) negative_keys = list(negative_counts.keys()) for i in range(len(positive_keys)): for j in range(i+1, len(positive_keys)): target = -(positive_keys[i] + positive_keys[j]) if target in negative_counts: count_triplets += positive_counts[positive_keys[i]] * \ positive_counts[positive_keys[j]] * negative_counts[target] for i in range(len(negative_keys)): for j in range(i+1, len(negative_keys)): target = -(negative_keys[i] + negative_keys[j]) if target in positive_counts: count_triplets += negative_counts[negative_keys[i]] * \ negative_counts[negative_keys[j]] * positive_counts[target] print(count_triplets + countDoubles(counts) + comb(counts[0], 3)) |