-
Notifications
You must be signed in to change notification settings - Fork 0
/
3sum with multiplicity.py
31 lines (27 loc) · 1.04 KB
/
3sum with multiplicity.py
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
class Solution(object):
def threeSumMulti(self, A, target):
MOD = 10**9 + 7
count = collections.Counter(A)
keys = sorted(count)
ans = 0
for i, x in enumerate(keys):
T = target - x
j, k = i, len(keys) - 1
while j <= k:
y, z = keys[j], keys[k]
if y + z < T:
j += 1
elif y + z > T:
k -= 1
else: # x+y+z == T, now calculate the size of the contribution
if i < j < k:
ans += count[x] * count[y] * count[z]
elif i == j < k:
ans += count[x] * (count[x] - 1) / 2 * count[z]
elif i < j == k:
ans += count[x] * count[y] * (count[y] - 1) / 2
else: # i == j == k
ans += count[x] * (count[x] - 1) * (count[x] - 2) / 6
j += 1
k -= 1
return int(ans % MOD)