Leetcode Three Sum

Question

From here:

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.

Note:

The solution set must not contain duplicate triplets.

This solution works and is pretty fast but not very memory efficient. I have some comments which I added when writing it to keep track of what I was doing. Specific questions are at bottom.

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        # Early exit conditions
        if len(nums) < 3:
            return []            

        results = {}
        vals = {}
        # Make easy to check if all target values exist
        for i, v in enumerate(nums):
            if v in vals:
                vals[v] += 1
            else:
                vals[v] = 1

        # Go through all matching numbers 
        # This works because each result is unique
        for i, num in enumerate(vals.keys()):

            # Modified two sum problem to get -i-j
            for j in list(vals.keys())[i:]:

                # If checking itself but only is in input once, can't work
                if num == j and vals[num] < 2:
                    continue

                target_val = -(num+j) 
                if target_val in vals:
                    valid = False
                    if vals[target_val] > 2:
                        valid = True
                    else:
                        # Check that this isn't either i/j (thus duplicate)
                        if num != target_val and j != target_val:
                            valid = True

                    if valid and tuple(sorted((num, j, target_val))) not in results: 
                        results[tuple(sorted((num, j, target_val)))] = True

        return [list(r) for r in results.keys()]

1. I dislike the way I'm sorting then converting to a tuple. Is there a better way to do this?
2. I feel like I can clean up some of the if logic but I'm not totally sure how is best.
3. I dislike the list(vals.keys()) approach - is there a better way to do what I'm doing here? I am not a huge fan of it
4. This one might be out of scope for this Code Review question but this doesn't seem to generalize very well if there were say 4 elements or 8 elements in the target (vs 3). I'd have to continue to add nested lists in order for that to work and the strategy for if checking here would get very messy.

Answer 1

For reference, what I think is the most straightforward solution, that is not fast enough since it does all combinations of three.

from itertools import combinations

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        return set(tuple(c) for c in combinations(sorted(nums), 3) if sum(c) == 0)

And a longer, but valid solution.

• Take positive and negative numbers separately.

• Add all combinations of two positive numbers where their negative sum is among the negative numbers. Repeat for negative numbers.

• For each number in both negative and positive, add -n, 0, n

• If there are three or more zeroes, add 0,0,0

from itertools import combinations

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        positive = sorted([n for n in nums if n > 0])
        posSet = set(positive)
        negative = sorted([n for n in nums if n < 0])
        negSet = set(negative)
        zeroes = nums.count(0)
        valid = set((a,b,0-a-b) for (a,b) in combinations(positive, 2) if 0-a-b in negSet)
        valid = valid.union(set((a,b,0-a-b) for (a,b) in combinations(negative, 2) if 0-a-b in posSet))
        if(zeroes > 0):
            valid = valid.union(set((-a,0,a) for a in positive if 0-a in negSet))
        if(zeroes > 2):
            valid.add((0,0,0))
        return valid

Surprisingly good results:

Runtime: 360 ms, faster than 98.52% of Python3 online submissions for 3Sum. Memory Usage: 16.4 MB, less than 97.14% of Python3 online submissions for 3Sum.

Answer 2

1. Please don't have white space at the end of lines. There is never a need for anything other than a new line at the end of a line.
2. You don't need to use enumerate(nums) to build vals - i is never used.
3. You can use collections.defaultdict(int) to remove the need for the if in the for when you make vals.
4. You can use collections.Counter to create vals.
5. You can use itertools.combinations_with_replacement to reduce your nested for loops into one. This has the same \$O\$, but may help in the future if you want to do four sum.
6. It is confusing to use i and j to not be related to the same thing. i is "index", where j is "num2".
7. Since valid is just a Boolean, you can compress the if-else to not need the if.

8. I believe your vaild check is erroneous. If j and target_val are the same but vals[j] == 1 then your code will happily still plop it out. I also think your output is erroneous on num == j and vals[num] == 2.

   You can make a collections.Counter of the values and then check if this counter is a subset of the global one.

9. I suggest making two functions:

   1. A generator function, one that produces values that add to the desired sum.
   2. A function to filter duplicate values and normalize output.

class Solution:
    def _three_sum(self, nums: List[int]) -> Iterator[Tuple[int, int, int]]:
        if len(nums) < 3:
            return

        vals = set(nums)
        for i, j in itertools.combinations_with_replacement(vals, 2):
            k = -(i+j)
            if k not in vals:
                continue
            yield i, j, k

    def threeSum(self, nums: List[int]) -> List[List[int]]:
        output = set()
        amounts = collections.Counter(nums)
        for result in self._three_sum(nums):
            if all(
                amount <= amounts[num]
                for num, amount in collections.Counter(result).items()
            ):
                output.add(tuple(sorted(result)))
        return [list(vals) for vals in output]

---

To answer your questions:

1. No.
2. If you change _three_sum to _n_sum then the threeSum function will work fine.
3. I'm not sure what you mean by it but, list(vals) is a Pythonic way to write it. Albeit possibly more confusing. It is preferred to just use itertools.

4. This is off-topic here, and no this would probably not be extendable to more than sum of three.

   I know of this answer on CS.SE that sets the foundation for an n-sum solution. Whilst the question clearly has a different goal, you should be able to dissect the answer and mold it to what you need.

Answer 3

I really liked @JollyJoker's answer but it was not written in Pythonic enough format for my taste. I added some changes that take his core logic and improve upon it:

from itertools import combinations

class Solution:
    def threeSum(self, nums: List[int]) -> List[List[int]]:
        positive = sorted(n for n in nums if n > 0)
        posSet = set(positive)
        negative = sorted(n for n in nums if n < 0)
        negSet = set(negative)
        zeroes = nums.count(0)
        valid = {(a, b, -a-b) for (a,b) in combinations(positive, 2) if -a-b in negSet}
        valid.update((a, b, -a-b) for (a,b) in combinations(negative, 2) if -a-b in posSet)
        if zeroes > 0:
            valid.update((-a, 0, a) for a in positive if -a in negSet)
            if zeroes > 2:
                valid.add((0, 0, 0))
        return valid

Results:

Runtime: 340 ms, faster than 98.68% of Python3 online submissions for 3Sum.

Memory Usage: 16.6 MB, less than 77.86% of Python3 online submissions for 3Sum.