Problem statement

Solution

An efficient solution to this problem can be easily derived from the solution to 167. Two Sum II - Input Array Is Sorted, so I highly recommmend solving that problem first.

This problem can be solved as a sequence of solutions to the aforementioned Two Sum II problem. Specifically, recall that in the Two Sum II problem, we are given an array numbers of integers that are sorted in non-decreasing order. We are then asked to find the two integers numbers[i] and numbers[j] such that numbers[i] + numbers[j] == target, where target is a given integer.

In the 3Sum problem, we are given an integer array nums, not necessarily sorted, and we are asked to return all the triplets [nums[i], nums[j], nums[k]] such that nums[i] + nums[j] + nums[k] == 0 and the triplets are unique (i.e. if (i1, j1, k1) != (i2, j2, k2), then [nums[i1], nums[j1], nums[k1]] != [nums[i2], nums[j2], nums[k2]]).

At a high level, we can solve the 3Sum problem as follows:

  1. Sort the nums array in non-decreasing order.
  2. Initialize an array results to an empty list to store the triplets.
  3. For i = 0, ..., len(nums) - 1:
    a. Set target to be -nums[i].
    b. Try to solve the Two Sum II problem with nums[l:r+1] and target as inputs.
    c. …
  4. Return results.

I intentionally left step 3(c) blank because it contains details that are important for the actual implementation but are not relevant for understanding the big picture.

Here is how this problem is actually solved in Python, including the aforementioned implementation details:

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

class Solution():
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()
        triplets = []
        for i in range(len(nums)):
            # If nums[i] > 0, then there does not exist triplets [[nums[i], nums[j], nums[k]]
            # such that nums[i] + nums[j] + nums[k] == 0 for every j > i and k > i. This
            # is because the nums array is sorted in non-decreasing order, and so
            # nums[j] >= nums[i] > 0 and nums[k] >= nums[i] > 0 for every j > i and k > i.
            if nums[i] > 0:
                break

            # Keep moving the pointer i so that duplicates are skipped, as required by
            # the problem.
            if i > 0 and nums[i] == nums[i - 1]:
                continue

            # Start solving a variant of Two Sum II
            l, r = i + 1, len(nums) - 1
            while l < r:
                threeSum = nums[i] + nums[l] + nums[r]
                if threeSum > 0:
                    r -= 1
                elif threeSum < 0:
                    l += 1
                else:
                    triplets.append([nums[i], nums[l], nums[r]])
                    # In Two Sum II, we return the pair [nums[l], nums[r]] at the end.
                    # However, because we are solving several Two Sum II problems
                    # consecutively, we have to move either the left pointer to the right
                    # or the right pointer to the left to keep avoid getting stuck in
                    # an infinite loop
                    l += 1
                    # Skip duplicates after moving the l pointer to the right
                    while nums[l] == nums[l - 1] and l < r:
                        l += 1
        return triplets