3-Sum Problem in Python

Question

I attempted the 3-Sum problem on Leetcode, where the problem asks to find all possible triplets of numbers in a given list such that their sum is 0. My code worked, but it exceeded the time limit for 2 of the 313 cases. I believe my solution is in \$O(n^2)\$, but I think there is a line in my code where I sort a triplet within a for-loop, which might change the time complexity.

Here is the problem statement from Leetcode:

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.

class Solution:
    def threeSumHelper(self, nums, target):
        D = {}
        n = len(nums)
        solutions = []
        for x in range(n):
            complement = target - nums[x]
            if complement in D.values():
                solutions.append([complement, nums[x]])
            D[x] = nums[x]
        return solutions

    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[List[int]]
        """
        D = {}
        n = len(nums)
        solutions = {}
        key = 0
        for x in range(n - 2):
            target = -(nums[x])
            pairs = self.threeSumHelper(nums[x+1:n], target)
            if not len(pairs) == 0:
                for pair in pairs:
                    pair.append(nums[x])
                    pair.sort()
                    if not (pair in solutions.values()):
                        if key in solutions:
                            key += 1
                        solutions[key] = pair
        return list(solutions.values())

The algorithm basically iteratively takes one number out of the list of numbers provided and sees whether there are two numbers in the remaining list that sum to the negative of the removed number (so that all three numbers sum to 0). This is done using a dictionary. Then, to avoid repetition, I sort the solutions and store them in a dictionary and return the list of values in the dictionary.

I just want to know if there is any way to make my code faster either by eliminating the need to sort each triplet or otherwise?

Answer 1

Looking at your code, this jumps out at me:

    D = {}

    for x in range(n):

        if complement in D.values():

The D variable here is a dictionary. It is designed for \$O(1)\$ lookup of a key. But you are looping over n items, performing a test (complement in D.values()) that is \$O(m)\$ itself, since D.values() is an iterable and the in test will be a linear search.

I suspect you can get a lot of performance back if you figure out how to not do that: it looks like you are treating D as a list, since you are using x as the key and nums[x] as the value.

I'll point out to you that unlike lists, tuples are immutable and therefore can be hashed, which makes them valid as dict keys. So you can use a dict of tuples. Or a set of tuples.