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
line1 = input()
n = int(line1)

line2 = input()
line2 = list(map(int, line2.split(" ")))

uc = line2

ps = [uc[0]]
for i in range(1, len(uc)):
    ps.append(ps[-1] + uc[i])

buc = []

for i in range(len(uc)):
    for j in range(i, len(uc)):

        subtract = 0
        if i > 0:
            subtract = ps[i - 1]
        buc.append(ps[j] - subtract)

from typing import List
def threeSum(nums: List[int]) -> List[List[int]]:
    nums.sort()
    triplets = 0

    for i in range(len(nums)):
        rest = -nums[i]
        lo = i + 1
        hi = len(nums) - 1

        while lo < hi:
            cur = nums[lo] + nums[hi]
            if cur < rest:
                lo += 1
            elif cur > rest:
                hi -= 1
            else:

                cnt1 = 1
                cnt2 = 1
                while lo + 1 < hi and nums[lo] == nums[lo + 1]:
                    lo += 1
                    cnt1 += 1
                while lo + 1 < hi and nums[hi] == nums[hi - 1]:
                    hi -= 1
                    cnt2 += 1

                if nums[lo] != nums[hi]:
                    triplets += cnt1*cnt2
                else:
                    triplets += (cnt1+cnt2)*(cnt1+cnt2-1) // 2
                lo += 1

    return triplets

print(threeSum(buc))