3Sum

The problem description can be found in LeetCode.

Approach

Before solving this problem be sure to understand:

1. Two Sum.
2. Two Sum II - Input Array Is Sorted.

The most efficient way to solve this problem is to sort the input array, and then use this property to efficiently look for the triplets that add up to 0.

In fact, we can "lock" the first element and then look for the other two in the remaining ordered list. In this way, since we are looking for:

$$a+b+c=0$$

By locking the first element, we can reduce our problem to:

$$b+c=-a$$

So, we can apply a Two Sum II in the ordered list in order to get two numbers ($b$ and $c$) that add up to $-a$. We can repeat this process until we have numbers in the list.

Lastly, we need to be careful about the edge case, which consists of not locking the same element more than once to avoid to have duplicate solutions. For example, with:

nums = [-1, -1, 0, 0, 1]

We can lock the first -1 (index 0) and obtain a solution with 0 and 1. However, we have to skip the next -1 (index 1) to avoid to include the same solution in the result again.

Solution

from typing import List


class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        nums.sort()

        prev = None
        ans: List[List[int]] = []

        for i, num in enumerate(nums):
            # skip already seen element
            if num == prev:
                continue
            # if current element `a` is > 0,
            # it means that the next two `b` and `c` are > 0 since the list is sorted
            #
            # hence, there won't be any triplet which can add up to 0
            if num > 0:
                break

            prev = num
            # a + b + c = 0
            # b + c = -a
            #
            # hence, look for -a
            ans.extend([num, j, k] for j, k in self.two_sum(nums, -num, i + 1))

        return ans

    def two_sum(self, nums: List[int], target: int, left: int) -> List[List[int]]:
        right = len(nums) - 1
        ans = []

        while left < right:
            psum = nums[left] + nums[right]
            if psum == target:
                ans.append([nums[left], nums[right]])

                left += 1
                right -= 1

                # move one of the two pointers to avoid duplicated pairs
                while left < len(nums) and nums[left - 1] == nums[left]:
                    left += 1
            elif psum < target:
                left += 1
            else:
                right -= 1

        return ans

Complexity

• Time: $O(n^2)$
• Space: $O(n)$

Where $n$ is the size of the nums list.