Three Sum

This problem is modified from the following leetcode link.

Leetcode link: https://leetcode.com/problems/3sum/

Question#

Given an array nums of n integers, are there elements a, b, c in nums such that a + b + c = target?

Find all triplets in the array which gives the sum of the target (you may assume that there are only unique triplets).

Given array nums = [-1, 1, 2, 3, 4, 5, 6, 7], target = 6.

A solution set is:
[
  [1, 2, 3],
  [-1, 1, 6]
]

Hint: Think how we can build upon our solution in the Two-Sum problem.

Method 1: Brute Force#

As with any problem, we always start off with the brute force technique. What is the simplest way to get the answer out, without any concern about the time complexity etc.?

We can simply use 3 for-loops and find out the indices like in our twoSum problem.

This will give us the right results, but such a method will be a poor solution for a huge array. Though space complexity is in constant time of O(1), with 3 for-loops we have a time complexity of O(n³).

Method 2: Reuse our twoSum function#

Does the problem seem familiar? The three sum problem is very much like the two sum, just that we are finding 3 numbers instead of 2. A great way to solve this problem is to use the twoSum method we wrote with slight modifications.

The twoSum method will now just use a set instead of a dictionary as we don't need to return the indices anymore.

def twoSum(nums: List[int], target: int) -> List[int]:
    num_set = set()
    for num in nums:
        complement = target - num
        if(complement in num_set):
            return [num, complement]
        else:
            num_set.add(num)
    return []

Let's write our threeSum function now.

from typing import List

def threeSum(nums: List[int], target) -> List[List[int]]:
    triplets = []
    for i in range(len(nums)):
        required = target - nums[i]
        if((i+1) < len(nums)):
            result = twoSum(nums[i+1:], required)
            if(result != []):
                x = [nums[i], result[0], result[1]]
                triplets.append(x)
    return triplets

print(threeSum([-1, 1, 2, 3, 4, 5, 6, 7], 6))
"""
Result:
[[-1, 4, 3], [1, 3, 2]]
"""

In each iteration, we calculate what is the required target we need to find a pair from. And we use our twoSum function to find that pair. If we do find a pair, we add it to our list of triplets.

Note that we are slicing the array to find the pair from using sums[i+1:]. This will search a subarray of the numbers as we keep moving forward in the iteration.

Finally, we return the triplets which has the list of triplets we found.

This yields a time complexity of O(n²) since for each element we have a subarray we need to search through for a pair. We will have a space complexity of O(n) from the twoSum method.