1
\$\begingroup\$

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))
\$\endgroup\$
  • 1
    \$\begingroup\$ You currently have a small error? typo? in your code: len_n, res, n_map. target = len(nums), set(), dict(), is not actually ok. I am not sure how this is supposed to look like, so please just correct it. \$\endgroup\$ – MegaIng Feb 12 at 22:03
  • \$\begingroup\$ thank you. updated already..it is typo when i post the code.. \$\endgroup\$ – sln Feb 13 at 1:21
2
\$\begingroup\$

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.

\$\endgroup\$
  • \$\begingroup\$ thank you. the review make so much sense now. However, couple of question for line ` for j, a, b in n_map[target - c]:` what if no key can found on n_map? and can you explain a littit bit about why it is good choice to use yield yield [c, a, b] here? \$\endgroup\$ – sln Feb 13 at 16:11
  • \$\begingroup\$ If there is no element with key target-c in n_map, this defaultdict returns an empty list. The for-loop starts iterating over it, but since it's empty, there are no iterations. On your second question, the order of c, a, b has no effect \$\endgroup\$ – Maarten Fabré Feb 13 at 16:28
  • \$\begingroup\$ understand, thank you.. great solution for your alternative approach as well \$\endgroup\$ – sln Feb 13 at 16:30

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.