Leetcode 3Sum problem using hashmap in Python

Question

For practice, I solved Leetcode 15. 3Sum question:

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.

My code pass all testcases from my location but get Time Limit Exceeded from leetcode. Can anyone suggest me how to improve my code?

def threeSum(nums):
    """
    :type nums: List[int]
    :rtype: List[List[int]]
    """
    if len(nums) < 3: return []
    len_n, res, n_map, target = len(nums), set(), dict(), 0
    nums.sort()
    for i, n in enumerate(nums):
        for j, c in enumerate(nums[i + 1:], i + 1):
            s = n + c
            n_map[s] = n_map.get(s, []) + [(i, j)]

    for i, n in enumerate(nums):
        s = target - n
        for k in n_map.get(s, []):
            if i > k[1]:
                res.add((nums[i], nums[k[0]], nums[k[1]]))
    return list(map(list, res))

Answer 1

I think the current way of solving this problem seems an acceptable, albeit naive, way to solve it, but there are some tweaks that can enhance the readability.

variables

len_n, res, n_map, target = len(nums), set(), dict(), 0

is both unclear and unnecessary.

len_n is never used, res is only used far further in the code

collections.defaultdict

You do a lot of n_map.get(s, []). Simpler would be to define n_map as a collectcions.defaultdict(list), and then for example just do n_map[s].append((i, j))

indices

You add (i, j) to n_map, only to later retrieve them as tuple k. It would be easier to use tuple unpacking:

for k, n in enumerate(nums): # i is used
    s = target - n
    for i, j in n_map[s]:
        if k > j:
            res.add((nums[k], nums[i], nums[j]))

Since you only use i and j here to retrieve a and b, why not save them in n_map in the first place?

n_map = defaultdict(list)
for i, a in enumerate(nums):
    for j, b in enumerate(nums[i + 1 :], i + 1):
        n_map[a + b].append((j, a, b))
res = set()
for k, c in enumerate(nums):
    for j, a, b in n_map[target - c]:
        result = c, a, b
        if k > j:
            ...

res and yield

Defining res as a set is a good choice. I think it is easier to only add the tuple to res if it is not present yet, and yield it at the same time, instead of returning list(map(list, res)) at the end

In total this gives:

def three_sum_maarten(nums, target=0):
    """
    :type nums: List[int]
    :rtype: List[List[int]]In total this gives
    """
    if len(nums) < 3:
        return []
    n_map = defaultdict(list)
    nums = sorted(nums)
    for i, a in enumerate(nums):
        for j, b in enumerate(nums[i + 1 :], i + 1):
            n_map[a + b].append((j, a, b))
    res = set()
    for k, c in enumerate(nums):
        for j, a, b in n_map[target - c]:
            result = c, a, b
            if k > j and result not in res:
                yield [c, a, b]
                res.add(result)

With this leetcode boilerplate:

class Solution:
    def threeSum(self, nums: 'List[int]') -> 'List[List[int]]':
        return list(three_sum_maarten(nums))

This passes all but one scenario. The scenario it fails is nums = [0] * 3000

To tackle this scenario, you can filter all numbers so only maximum 3 of each are present in nums. I do this with the help of a collections.Counter:

def prepare_nums(nums):
    counter = Counter(nums)

    for n, c in sorted(counter.items()):
        yield from [n] * min(c, 3)

and then nums = list(prepare_nums(nums)) instead of nums = sorted(nums)

---

Alternative approach

You make about half of all combinations of 2 numbers in nums. One extra bit of knowledge you can use to reduce this is to take into account that at least 1 negative and 1 positive number need to be present in each triplet.

counter = Counter(nums)
positives = [i for i in counter if i > 0]
negatives = [i for i in counter if i < 0]

for a, b in product(positives, negatives):
    c = -(a + b)
    if c not in counter:
        continue
    result = a, b, c

and then only yield the correct, unique results

    result = a, b, c
    if c == a:
        if counter[a] >= 2:
            yield result
    elif c == b:
        if counter[b] >= 2:
            yield result
    elif a > c > b:
        yield result

and yield 1 (0, 0, 0) triplet if there are 3 or more 0s present

if counter[0] >= 3:
    yield (0, 0, 0)

This solution is about 10 times faster, and uses 30 times less memory.