Leetcode Three Sum in Python

Question

Here's my solution for the Leet Code's Three Sum problem -- would love feedback on (1) code efficiency and (2) style/formatting. This is Python 3.

Problem:

Given an array nums of n integers, are there elements a, b, c in nums 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.

Example:

Given array nums = [-1, 0, 1, 2, -1, -4],

A solution set is: [ [-1, 0, 1], [-1, -1, 2] ]

My solution:

def threeSum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """

    def _twoSum(nums_small, target):
        '''
        INPUT: List of numeric values, int target 
        OUTPUT: list: pair(s) from nums_small that add up to target and -1*target
        '''
        nums_dict = {}
        return_lst = []

        for indx, val in enumerate(nums_small, target):
            if target - val in nums_dict: 
                return_lst.append([val, target - val, -1*target]) 
            nums_dict[val] = indx
        return return_lst 

    return_dict = {}

    for indx, val in enumerate(nums):
        for solution in _twoSum(nums[:indx] + nums[indx+1:], -1*val):
            if sorted(solution) not in return_dict.values():
                return_dict[str(indx)+str(solution)] = sorted(solution) # ... hacky way of storing multiple solutions to one index ...

    return list(return_dict.values())

Answer 1

Code

"When all you have is a hammer, everything looks like a nail."

Your tool chest seems to have dict, which you are using to implement a container which doesn't contain any duplicates. There is a name for that tool. It is called as set. Study it; you will find it very useful.

Consider:

    nums_dict = {}

    for indx, val in enumerate(nums_small, target):
        if target - val in nums_dict: 
            pass 
        nums_dict[val] = indx

You are never using the value indx; you are just storing it and never retrieving it. You could just as easily write:

    nums_dict = {}

    for val in nums_small:
        if target - val in nums_dict: 
            pass 
        nums_dict[val] = True

which avoids the enumerate(), with the odd initial value, and the indices which did nothing. But this is still storing the values True, and never retrieving them. You don't need to do this, if you use a set:

    nums_dict = set()

    for val in nums_small:
        if target - val in nums_dict: 
            pass 
        nums_dict.add(val)

You again use a dict for return_dict to ensure you don't have any duplicate solutions. Actually, this code is a worse that the former, because you are using in return_dict.values(), so instead of an \$O(1)\$ key lookup, you're doing an \$O(n)\$ search through the entire list of values! If that wasn't bad enough, you are sorting solutions twice: once to see if they already exist in return_dict, and a second time to add it into the dictionary if it wasn't found. Saving the sorted solution would avoid the second sort.

You could replace return_dict with a set as well, but there is a catch, you can't add a list into a set, because the items in a set must be hashable to a constant value but lists are unhashable (because they can be mutated). You can get around this by using tuples, which can't be mutated, so they can be hashed.

 return_dict = set()

 for index, val in ...:
    for solution in ...:
        sorted_solution = tuple(sorted(solution))
        return_dict.add(sorted_solution)

 return list(return_dict)

The above returns a list of tuples. If you need a list of lists, that can be easily obtained as well, using list comprehension ...

 return [ list(solution) for solution in return_dict ]

... which is another useful tool for your tool chest.

---

Algorithm

Using _twoSum() to solve the threeSum() problem is good reuse of code. Kudos.

However, looping over all indices and for each index, constructing a new list omitting just that one value from the list nums[:indx] + nums[indx+1:] is extremely inefficient.

Since "Leetcode" is a puzzle solving site, I won't give you the better algorithm, but I will say that splitting nums into len(nums) different lists is not the way to go. I would split it into 2 lists, or maybe 3 lists tops. Good luck!