algorithms/array/three_sum.py at master · mindreframer/algorithms · GitHub

History

46 lines (39 loc) · 1.11 KB

Raw

"""

Given an array S of n integers, are there elements a, b, c in S

such that a + b + c = 0?

Find all unique triplets in the array which gives the sum of zero.

Note: The solution set must not contain duplicate triplets.

For example, given array S = [-1, 0, 1, 2, -1, -4],

A solution set is:

[

[-1, 0, 1],

[-1, -1, 2]

]

"""

def three_sum(nums:"List[int]")->"List[int]":

res = []

nums.sort()

for i in range(len(nums)-2):

if i > 0 and nums[i] == nums[i-1]:

continue

l, r = i+1, len(nums)-1

while l < r:

s = nums[i] + nums[l] + nums[r]

if s > 0:

r -= 1

elif s < 0:

l += 1

else:

# found three sum

res.append((nums[i], nums[l], nums[r]))

# remove duplicates

while l < r and nums[l] == nums[l+1]:

l+=1

while l < r and nums[r] == nums[r-1]:

r -= 1

l += 1

r -= 1

return res

if __name__ == "__main__":

x = [-1,0,1,2,-1,-4]

print(three_sum(x))

algorithms/array/three_sum.py at master · mindreframer/algorithms · GitHub

Expand file tree

Breadcrumbs

three_sum.py

Latest commit

History

Breadcrumbs

three_sum.py

File metadata and controls