Optimizing solution to Three Sum

Question

I am trying to solve the 3 Sum problem stated as:

Given an array S of n integers, are there elements a, b, c in S such that a + b + c = 0? Find all unique triplets in the array which gives the sum of zero.

Note: The solution set must not contain duplicate triplets.

Here is my solution to this problem:

def threeSum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    nums.sort()
    n = len(nums)
    solutions = []
    for i, num in enumerate(nums):
        if i > n - 3:
            break

        left, right = i+1, n-1
        while left < right:
            s = num + nums[left] + nums[right] # check if current sum is 0
            if s == 0:
                new_solution = [num, nums[left], nums[right]]
                # add to the solution set only if this triplet is unique
                
                if new_solution not in solutions: 
                    solutions.append(new_solution)
                right -= 1
                left += 1
            elif s > 0:
                right -= 1
            else:
                left += 1

    return solutions

This solution works fine with a time complexity of O(n**2 + k) and space complexity of O(k) where n is the size of the input array and k is the number of solutions.

While running this code on LeetCode, I am getting TimeOut error for arrays of large size. I would like to know how can I further optimize my code to pass the judge.

P.S: I have read the discussion in this related question. This did not help me resolve the issue.

Answer 1

A couple of improvements you can make to your algorithm:

1) Use sets instead of a list for your solution. Using a set will insure that you don't have any duplicate and you don't have to do a if new_solution not in solutions: check.

2) Add an edge case check for an all zero list. Not too much overhead but saves a HUGE amount of time for some cases.

3) Change enumerate to a second while. It is a little faster. Weirdly enough I am getting better performance in the test with a while loop then a n_max = n -2; for i in range(0, n_max): Reading this question and answer for xrange or range should be faster.

NOTE: If I run the test 5 times I won't get the same time for any of them. All my test are +-100 ms. So take some of the small optimizations with a grain of salt. They might NOT really be faster for all python programs. They might only be faster for the exact hardware/software config the tests are running on.

ALSO: If you remove all the comments from the code it is a LOT faster HAHAHAH like 300ms faster. Just a funny side effect of however the tests are being run.

I have put in the O() notation into all of the parts of your code that take a lot of time.

def threeSum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    # timsort: O(nlogn)
    nums.sort()

    # Stored val: Really fast
    n = len(nums)

    # Memory alloc: Fast
    solutions = []

    # O(n) for enumerate
    for i, num in enumerate(nums):
        if i > n - 3:
            break

        left, right = i+1, n-1

        # O(1/2k) where k is n-i? Not 100% sure about this one
        while left < right:
            s = num + nums[left] + nums[right]  # check if current sum is 0
            if s == 0:
                new_solution = [num, nums[left], nums[right]]
                # add to the solution set only if this triplet is unique

                # O(n) for not in
                if new_solution not in solutions:
                    solutions.append(new_solution)
                right -= 1
                left += 1
            elif s > 0:
                right -= 1
            else:
                left += 1

    return solutions

Here is some code that won't time out and is fast(ish). It also hints at a way to make the algorithm WAY faster (Use sets more ;) )

class Solution(object):
def threeSum(self, nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    # timsort: O(nlogn)
    nums.sort()

    # Stored val: Really fast
    n = len(nums)

    # Hash table
    solutions = set()

    # O(n): hash tables are really fast :)
    unique_set = set(nums)
    # covers a lot of edge cases with 2 memory lookups and 1 hash so it's worth the time
    if len(unique_set) == 1 and 0 in unique_set and len(nums) > 2:
        return [[0, 0, 0]]

    # O(n) but a little faster than enumerate.
    i = 0
    while i < n - 2:
        num = nums[i]

        left = i + 1
        right = n - 1

        # O(1/2k) where k is n-i? Not 100% sure about this one
        while left < right:
            # I think its worth the memory alloc for the vars to not have to hit the list index twice. Not sure
            # how much faster it really is. Might save two lookups per cycle. 
            left_num = nums[left]
            right_num = nums[right]
            s = num + left_num + right_num  # check if current sum is 0
            if s == 0:
                # add to the solution set only if this triplet is unique
                # Hash lookup
                solutions.add(tuple([right_num, num, left_num]))
                right -= 1
                left += 1
            elif s > 0:
                right -= 1
            else:
                left += 1
        i += 1

    return list(solutions)