Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

3SUM (finding all unique triplets in a list that equal 0)

Tags:

python

list

I am working on the 3SUM problem (taken from leetcode), which takes a list as input and finds all unique triplets in the lists such that a+b+c=0. I am not really sure what my code is doing wrong, but it currently returns an empty list for this list [-1, 0, 1, 2, -1, -4], so it is not recognizing any triplets that sum to 0. I would appreciate any suggestions or improved code.

Here's my code:

result = []
nums.sort()
l = 0
r=len(nums)-1
for i in range(len(nums)-2):
    while (l < r):
        sum = nums[i] + nums[l] + nums[r]
        if (sum < 0):
            l = l + 1
        if (sum > 0):
            r = r - 1
        if (sum == 0): 
            result.append([nums[i],nums[l],nums[r]])
print(result)
like image 780
Jane Sully Avatar asked Sep 19 '25 22:09

Jane Sully


1 Answers

A couple things to note.

  1. Don't use sum as a variable name because that's a built-in function.
  2. Your indexing is a little problematic since you initialize l = 0 and have i begin at 0 as well.
  3. Don't rest on your laurels: increment the value of l when you find a successful combination. It's really easy to forget this step!

An edited version of your code below.

nums = [-1, 0, 1, 2, -1, -4]
result = []
nums.sort()
r=len(nums)-1
for i in range(len(nums)-2):
    l = i + 1  # we don't want l and i to be the same value.
               # for each value of i, l starts one greater
               # and increments from there.
    while (l < r):
        sum_ = nums[i] + nums[l] + nums[r]
        if (sum_ < 0):
            l = l + 1
        if (sum_ > 0):
            r = r - 1
        if not sum_:  # 0 is False in a boolean context
            result.append([nums[i],nums[l],nums[r]])
            l = l + 1  # increment l when we find a combination that works

>>> result
[[-1, -1, 2], [-1, 0, 1], [-1, 0, 1]]

If you wish, you can omit the repeats from the list.

unique_lst = []
[unique_lst.append(sublst) for sublst in result if not unique_lst.count(sublst)]

>>> unique_lst
[[-1, -1, 2], [-1, 0, 1]]

Another approach uses itertools.combinations. This doesn't require a sorted list.

from itertools import combinations

result = []
for lst in itertools.combinations(nums, 3):
    if sum(lst) == 0:
        result.append(lst)

A nested for loop version. Not a big fan of this approach, but it's basically the brute-force version of the itertools.combinations solution. Since it's the same as approach as above, no sort is needed.

result = []
for i in range(0, len(nums)-2):
    for j in range(i + 1, len(nums)-1):
        for k in range(j + 1, len(nums)):
            if not sum([nums[i], nums[j], nums[k]]):  # 0 is False
                result.append([nums[i], nums[j], nums[k]])
like image 83
3novak Avatar answered Sep 21 '25 13:09

3novak