This is the first in a series of posts where I will walk through a problem. These are not intended to be the most efficient solutions, but rather a way to show my thought process and how I approach problems. I'm starting with 3Sum because it's a great follow up to the Two Sum problem, which is most people's introduction to LeetCode.

Problem statement

Given an integer array nums, return all the triplets [nums[i], nums[j], nums[k]] such that i != j, i != k, and j != k, and nums[i] + nums[j] + nums[k] == 0.

Notice that the solution set must not contain duplicate triplets.

Example 1:

Input: nums = [-1,0,1,2,-1,-4] Output: [[-1,-1,2],[-1,0,1]]

Example 2:

Input: nums = [0,0,0] Output: [[0,0,0]]

Naive/Brute force solution

The naive solution is to use three nested loops to check every possible combination of three numbers. This is a O(n^3) solution, which is not great. It's very basic so I'm not doing to code it here. We will do something better.

Better solution

This problem can be rewritten in the style of the Two Sum problem. Rather than trying to find nums[i] + nums[j] + nums[k] == 0, we can instead look for nums[i] + nums[j] == -nums[k]. This looks more like the Two Sum problem, where -nums[k] is the target number. We can then use the same approach as the Two Sum problem, but we need to make a few changes.

We will start by going through the array. Each value we look at becomes our new target. We then perform a Two Sum search on the rest of the array, looking for two numbers that add up to our target. If we find a match, we add the three numbers to our result set. We then move on to the next number in the array and repeat the process.

Before we get onto programming this, it's always good to think about the time (and space) complexity of our solution. This also makes sure that we actually understand what it is that we'll be coding. We will be going through the array once, and for each number we will be performing a Two Sum search on the rest of the array. This means that we will be going through the array n times, and then perform the Two Sum search, which we can do in O(n). This gives us a time complexity of O(n^2).

Missing piece

If we began to code this, we will run into an issue because of the rule that "the solution set must not contain duplicate triplets". We can see this being an issue in the example of [-1,0,1,2,-1,-4], since [-1, 0, -1] add to 0, but so is [0, 1, -1]. We want to count these are the same solution. The easiest way to deal with this is to simply sort this array before adding it to the set. This doesn't make our time complexity any worse, since sorting is O(n log n), which is less than O(n^2). More important than however is that all our answers are triplets, meaning it is always an array of size 3, so the sorting is actually in O(1)/constant time for this problem. It doesn't matter how big the array is, we will always be sorting 3 numbers.

Can we do better?

Yes we can! Although the improvement is also O(n^2) time complexity, it's a nicer solution and can be done without the use of a hash table like we do with traditional Two Sum. Since we require sorting to avoid duplicates anyway, and sorting is quicker than O(n^2), we can sort the entire array which lets us use the two pointer method to solve the Two Sum portion of this problem. (You can learn more about that by trying out Two Sum with a sorted input.)

This solution looks like this:

1. Sort the input array
2. Loop through the values in the array and set the current value as the target. If we've already used this element as a target, skip it. Because our input array is sorted, we can check if the current value is the same as the previous value and just move forward. If we tried to do the Two Sum search on the same value, we would just get the same result anyway!
3. Perform the Two Sum search using the pointer method with the target value from step 2. If we find a match, add the three numbers to our result set.

A brief overview of the pointer method:

1. Set two pointers, one at the start of the sorted array and one at the end.
2. If the sum of the two pointers is less than the target, move the left pointer forward. This increases the sum.
3. If the sum of the two pointers is greater than the target, move the right pointer back. This reduces the sum.
4. If the sum of the two pointers is equal to the target, we have a match! Add the two numbers to our result set.
5. Move both pointers forward/backward and repeat the process until the pointers meet.

Code

def threeSum(nums: List[int]) -> List[List[int]]:
    res = []
    nums.sort()

    for i, n in enumerate(nums):
        # skip this number since we don't want to get the same results from the Two Sum search
        if i > 0 and n == nums[i - 1]:
            continue

        # from re-arranging the equation to make it Two Sum
        target = -n
        
        # Two Sum, two pointer solution, which is applicable because array is sorted
        l, r = i + 1, len(nums) - 1
        while l < r:
            if nums[l] + nums[r] < target:
                # increase sum
                l += 1
            elif nums[l] + nums[r] > target:
                # reduce sum
                r -= 1
            else:
                # when nums[l] + nums[r] == target
                # note that we add n to the res array, not target
                res.append([n, nums[l], nums[r]])
                # move pointers forward/backward
                l += 1
                r -= 1
                # slight edge case to note:
                # to avoid duplicates, we also want to move l forward if the next position is the same
                while l < len(nums) and nums[l] == nums[l - 1]:
                    l += 1

    return res

References

• 3Sum on LeetCode
• Two Sum on LeetCode
• Two Sum with sorted input array on LeetCode