Interview challenge · medium · array, two-pointers, sorting
3Sum
El reto
El enunciado va en inglés, igual que en el vídeo.
Given an integer array, return all the unique triplets [a, b, c] whose sum is zero. The output must not contain the same triplet twice.
Example
[-1, 0, 1, 2, -1, -4] -> [[-1, -1, 2], [-1, 0, 1]]
[0, 0, 0, 0] -> [[0, 0, 0]]Idea: sort the array, then fix the first number and look for the other two with two pointers that close in from both ends: if the sum is too small, move the left one; if it is too big, move the right one; on zero, keep the triplet and move both. Skipping repeated values keeps the triplets unique. Time O(n²), extra space O(1) besides the output.
La solución
# 3Sum - ayotl.dev
def three_sum(nums):
a, out = sorted(nums), []
for i in range(len(a) - 2):
# skip a repeated first number
if i and a[i] == a[i - 1]:
continue
lo, hi = i + 1, len(a) - 1
# two pointers close in on zero
while lo < hi:
s = a[i] + a[lo] + a[hi]
if s == 0:
t = [a[i], a[lo], a[hi]]
if t not in out:
out.append(t)
if s <= 0:
lo += 1
if s >= 0:
hi -= 1
return out
print(three_sum([-1, 0, 1, 2, -1, -4]))
print(three_sum([0, 0, 0, 0]))
$ python3 three_sum.py
[[-1, -1, 2], [-1, 0, 1]]
[[0, 0, 0]]