Python Programming Challenge 12: Finding three numbers that sum to zero

In today’s post, we’ll be working on another challenge from Leetcode: https://leetcode.com/problems/3sum/.

For this challenge, we need to write a function called threeSum() that accepts a list and returns all the unique triplets in the list that gives a sum of zero.

For instance, suppose the list is [-1, 0, 1, 2, -1, -4], the function should give the following output:

[[-1, -1, 2], [0, -1, 1]]

Suggested Approach

One way to complete this challenge is to use three “pointers” to store the locations (i.e. indexes) of three different elements in the list.

Let’s call the three pointers main, left and right. Suppose the list is called nums and stores the list [-1, -1, 1, 2, -1, 1, 2, 3].

We need to sort this list first. After sorting, nums becomes [-1, -1, -1, 1, 1, 2, 2, 3].

At the beginning, we’ll let main store the index of the first element in nums (i.e. main = 0), left store the second (left = 1) and right store the last (right = 7, since the last element has an index of 7).

In other words,

nums[main] = -1
nums[left] = -1
nums[right] = 3

We first check the value of nums[main] + nums[left] + nums[right].

If the sum is greater than zero, we need to reduce the sum. Since right is currently pointing at the largest number in the list, we move right one step to the left.

If the sum is smaller than zero, we need to increase the sum. Since left is currently pointing at the smallest number in the list (other than nums[main]), we’ll move left one step to the right.

Each time we find a case where the sum is zero, we append the three numbers as a list to the result and move left and right one unit towards the center.

We keep doing this until the left pointer is no longer smaller than the right pointer.

When that happens, we need to update the value of main.

We move main one unit to the right, reset the left and right pointers and repeat the process above.

left always starts one unit to the right of main and right always starts at the end of the list.

The diagram below illustrates how this works.

In the diagram above, we see how this technique works for the first eight steps. We need to repeat the process until main, left and right point at the last three elements in the list. That way, we would have considered all the possible combinations of main, left and right.

However, if you study the diagram above, you may notice something amiss. In steps 2, 3 and 6, we append the same list [-1, -1, 2] to the result. This violates the requirement that the function should only return unique triplets that give a sum of zero.

Preventing duplicates is the hardest part of this question. Consider the following cases where duplicates can occur:

Case 1:

nums = [-2, -2, 0, 2]

When main = 0, left = 2, right = 3, we have the triplet [-2, 0, 2], which is a valid answer.

However, when main = 1, left = 2, right = 3, we have the same triplet.

Case 2:

nums = [-5, 0, 1, 1, 2, 3, 4, 4]

When main = 0, left = 2, right = 7, we have the triplet [-5, 1, 4] which is a valid triplet.

However, when main = 0, left = 3, right = 6, we have the same triplet.

How can we prevent duplicates like the 2 cases above from happening?

One way is to ensure that each time we update main, left or right, the element at the new position of these pointers should have a different value from the element at the original position.

For instance, in case 1 above, when we update main (from main = 0), we should update it to a value of 2 (instead of 1), since main[2] is a new value while main[1] is the same value as main[0].

The diagram below shows a visual representation of how main, left and right can be updated.

Based on the description above, try working on the threeSum() function yourself.

Suggested Solution

def threeSum(nums): 
    nums.sort()       
    l = len(nums)
    result = []

    main = 0
    left = 1
    right = l - 1

    # This outer while loop is for updating main
    while main < l - 2:

        # exit function if zero is impossible
        if nums[main] > 0:
            return result

        # this inner while loop is for updating left and right
        while left < right:  
            if nums[main] + nums[left] + nums[right] > 0:
                right -= 1
            elif nums[main] + nums[left] + nums[right] < 0:
                left += 1
            else:
                result.append([nums[left], nums[main], nums[right]])
                # this while loop is to prevent duplicates because of left
                while left < right and nums[left] == nums[left+1]:
                  left +=1
                # this while loop is to prevent duplicates because of right
                while left < right and nums[right-1] == nums[right]:
                  right -= 1

                left += 1
                right -= 1

        # this while loop is to prevent duplicates because of main
        while main >= 0 and main < l - 1 and nums[main] == nums[main+1]:
            main += 1  
                
        main += 1
        left = main + 1
        right = l - 1
        
    return result