[Problem Solving] Three Sum – Using Two Sum

To understand the following problem, we will be using what we learnt in the two sum problem. If you haven’t read it yet, do make sure to understand it before proceeding. Here’s the link.

Problem Statement

We have an array of integers called nums. We need to return three numbers let’s call it nums[i] = x, nums[j] = y and nums[k] = z that add up to 0. The indices of the returned numbers should all be unique ie., i != j != k

Assumptions

• Solution set should not contain duplicate triplets
• Empty array results in an empty set
• If there are less than three numbers in the input array, return an empty set

Examples

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

Output: [[-1,-1,2],[-1,0,1]]

Thought process

The brute force way to do this would be to loop using three for loops to search for three numbers in the array that sum up to 0. But this would have an O(n³) time complexity. So we need to be slightly better than that. The next best in time is an O(n²) solution. Remember that we should not return duplicate indices.

Let’s say nums[i] + nums[j] + nums[k] = 0

Then nums[i] + nums[j] = – nums[k]

This now resembles a two sum problem where we need to look for nums[i] and nums[j] such that they add up to a target = – nums[k]

To condense the problem, what we can do is loop through the array from k = 0 to k = n-1 and for every nums[k], find two numbers nums[i] and nums[j] that add up to a target = – nums[k].

How can we avoid adding duplicate triplets to the results list? A naive solution is to add all the duplicates to the result list and in the end return the set(result) to remove duplicates. This way we still do unnecessary work on finding the same two numbers that add up to a target value -nums[k].

To avoid executing a two sum loop on duplicates, let’s reuse our solution from the two sum problem where the input array was sorted and contained duplicates. We can sort the nums list and skip searching for two sum values whenever the current nums[k] = nums[k-1]. Sorting makes it easier to club duplicate values together. For ex:

Input: [-1,0,1,2,-1,-4]

Sorted List: [-4, -1, -1, 0 ,1,2]

One of the solutions here that adds up to 0 will be [-1, 0,1]. Now if we did not skip for target = nums[2] = -1, we will again get a solution set of [-1, 0, 1] in the result.

Complexity

O(N²) time complexity and O(1) extra space

Solution

def three_sum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    # Edge case when the array is empty
    if not nums:
       return []
    
    n = len(nums)
    # Another edge case when the input array does not have 3 elements
    if len(nums) < 3:
       return []
        
    results = []

    # Sorting helps to club duplicate values together
    nums.sort()
        
    for k in range(0, n):
        # If the current num is equal to the previous value, then skip execution because we will be getting a duplicate solution set
        if k > 0 and nums[k] == nums[k-1]:
            continue
                
            target = -1 * nums[k]
            start = k + 1
            end = n - 1
            
            two_sum(nums, target, results, start, end)
    return results

# Same as the two sum solution that we coded in the previous article
def two_sum(nums, target, results, s, e): 

    while s < e:
            
        if nums[s] + nums[e] == target:
            results.append([nums[s], nums[e], -target])
            s += 1 # Advance the starting pointer to continue to look for other solutions

            # Skip duplicate solutions in two sum
            while s < e and nums[s] == nums[s-1]:
                s += 1
            
        elif nums[s] + nums[e] < target:
           s += 1
        else:
           e -= 1