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 ofn_map.get(s, []). Simpler would be to definen_map as acollectcions.defaultdict(list), and then for example just don_map[s].append((i, j))

indices

You add(i, j) ton_map, only to later retrieve them as tuplek. 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 usei andj here to retrievea andb, why not save them inn_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 andyield

Definingres as a set is a good choice. I think it is easier to only add the tuple tores if it is not present yet, andyield it at the same time, instead of returninglist(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 isnums = [0] * 3000

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

def prepare_nums(nums):    counter = Counter(nums)    for n, c in sorted(counter.items()):        yield from [n] * min(c, 3)

and thennums = list(prepare_nums(nums)) instead ofnums = sorted(nums)

Alternative approach

You make about half of all combinations of 2 numbers innums. 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 more0s present

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

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

• \$\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 yieldyield [c, a, b] here?\$\endgroup\$

  sln

  – sln

  2019-02-13 16:11:01 +00:00

  CommentedFeb 13, 2019 at 16:11
• \$\begingroup\$If there is no element with keytarget-c inn_map, thisdefaultdict returns an empty list. Thefor-loop starts iterating over it, but since it's empty, there are no iterations. On your second question, the order ofc, a, b has no effect\$\endgroup\$

  Maarten Fabré

  – Maarten Fabré

  2019-02-13 16:28:06 +00:00

  CommentedFeb 13, 2019 at 16:28
• \$\begingroup\$understand, thank you.. great solution for your alternative approach as well\$\endgroup\$

  sln

  – sln

  2019-02-13 16:30:37 +00:00

  CommentedFeb 13, 2019 at 16:30

• python
• python-3.x
• programming-challenge
• time-limit-exceeded
• k-sum