Interview challenge · medium · array, two-pointers, sorting

3Sum

The challenge

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.

The solution

1# 3Sum - ayotl.dev
2def three_sum(nums):
3    a, out = sorted(nums), []
4    for i in range(len(a) - 2):
5        # skip a repeated first number
6        if i and a[i] == a[i - 1]:
7            continue
8        lo, hi = i + 1, len(a) - 1
9        # two pointers close in on zero
10        while lo < hi:
11            s = a[i] + a[lo] + a[hi]
12            if s == 0:
13                t = [a[i], a[lo], a[hi]]
14                if t not in out:
15                    out.append(t)
16            if s <= 0:
17                lo += 1
18            if s >= 0:
19                hi -= 1
20    return out
21
22
23print(three_sum([-1, 0, 1, 2, -1, -4]))
24print(three_sum([0, 0, 0, 0]))

$ python3 three_sum.py

[[-1, -1, 2], [-1, 0, 1]]

[[0, 0, 0]]

The video

Watch on YouTube →