“Time Limit Exceeded” on LeetCode's three-sum problem?

Question

I'm trying to solve the three-sum problem on LeetCode and I believe I've come up with some O ( n ^2) submissions, but I keep on getting a "Time Limit Exceeded" error.

For example, this solution using itertools.combinations :

from itertools import combinations

class Solution:
    def threeSum(self, nums):
        results = [triplet for triplet in combinations(nums, 3) if sum(triplet) == 0]
        return [list(result) for result in set([tuple(sorted(res)) for res in results])]

Results in the following error:

Similarly, this solution,

from itertools import combinations

class Solution:
    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[List[int]]
        """
        _map = self.get_map(nums)

        results = set()
        for i, j in combinations(range(len(nums)), 2):
            target = -nums[i] - nums[j]
            if target in _map and _map[target] and _map[target] - set([i, j]):
                results.add(tuple(sorted([target, nums[i], nums[j]])))
        return [list(result) for result in results]

    @staticmethod
    def get_map(nums):
        _map = {}
        for index, num in enumerate(nums):
            if num in _map:
                _map[num].add(index)
            else:
                _map[num] = set([index])
        return _map 

yields a "Time Limit Exceeded" for an input consisting of a long array of zeros:

This question has been asked before ( Optimizing solution to Three Sum ), but I'm looking for suggestions pertaining to these solutions specifically. Any idea what is making the solutions 'too slow' for LeetCode?

Update

It occurred to me that determined _map[target] - set([i, j]) - that is, whether the current set of indices are not also indices of the target value - could be expensive, so I should first look up whether the corresponding number pair has been seen or not. So I tried this:

from itertools import combinations

class Solution:
    def threeSum(self, nums):
        """
        :type nums: List[int]
        :rtype: List[List[int]]
        """
        _map = self.get_map(nums)

        results = set()
        seen = set()
        for i, j in combinations(range(len(nums)), 2):
            target = -nums[i] - nums[j]
            pair = tuple(sorted([nums[i], nums[j]]))
            if target in _map and pair not in seen and _map[target] - set([i, j]):
                seen.add(pair)
                results.add(tuple(sorted([target, nums[i], nums[j]])))
        return [list(result) for result in results]

    @staticmethod
    def get_map(nums):
        _map = {}
        for index, num in enumerate(nums):
            if num in _map:
                _map[num].add(index)
            else:
                _map[num] = set([index])
        return _map

However, this fails on another test case with large input numbers:

Answer 1

This has worked for me, uses a few optimizations for a lot of repeated elements. We store the count of the appearances of each element and then only iterate over each different element. The rest is similar to what you have already done

from collections import Counter
import itertools

class Solution:
    def threeSum(self, nums):
        counts = Counter(nums)
        numSet = list(set(nums))
        result = set()

        for idx, num1 in enumerate(numSet):
            for idx2, num2 in enumerate(itertools.islice(numSet, idx, None), start=idx):
                num3 = (num1 + num2) * -1
                if num3 in counts:
                    d = Counter([num1, num2, num3])
                    if not any(d[n] > counts[n] for n in d):
                        result.add(tuple(sorted([num1, num2, num3])))

        return list(result)