CODEWITHSUNDEEP
pythonpython-3.x
Newmsy
Newmsy

Reputation: 13

Iterating multiple variables over a list without double counting?

I was working on part of a program in which I'm trying to input a list of numbers and return all groups of 3 numbers which sum to 0, without double or triple counting each number. Here's where I'm up to:

def threeSumZero2(array):
    sums = []
    apnd=[sorted([x,y,z]) for x in array for y in array for z in array if x+y+z==0]
    for sets in apnd:
        if sets not in sums:
            sums.append(sets)
    return sums

Is there any code I can put in the third line to make sure I don't return [0,0,0] as an answer.

This is my test list:

[-1,0,1,2,-1,4]

Thank you

*Edit: I should have clarified for repeated input values: the result expected for this test list is: 

[[-1,-1,2],[-1,0,1]]

Upvotes: 0

Views: 241

Answers (2)

Olivier Melançon
Olivier Melançon

Reputation: 22314

You want combinations without replacement, this is something offered by itertools. Your sums can then be made a set to remove the duplicates with regard to ordering.

from itertools import combinations

def threeSumZero2(array):
    sums = set()
    for comb in combinations(array, 3):
        if sum(comb) == 0:
            sums.add(tuple(sorted(comb)))
    return sums

print(threeSumZero2([-1,0,1,2,-1,4]))

Output

{(-1, -1, 2), (-1, 0, 1)}

This solution can also be written more concisely using a set-comprehension.

def threeSumZero2(nums):
    return {tuple(sorted(comb)) for comb in combinations(nums, 3) if sum(comb) == 0}

More efficient algorithm

Although, the above algorithm requires traversing all combinations of three items, which makes it O(n3).

A general strategy used for this kind of n-sum problem is to traverse the n-1 combinations and hash their sums, allowing to efficiently test them against the numbers in the list.

The algorithm complexity drops by one order of magnitude, making it O(n2)

from itertools import combinations

def threeSumZero2(nums, r=3):
    two_sums = {}
    for (i_x, x), (i_y, y) in combinations(enumerate(nums), r - 1):
        two_sums.setdefault(x + y, []).append((i_x, i_y))

    sums = set()
    for i, n in enumerate(nums):
        if -n in two_sums:
            sums |= {tuple(sorted([nums[idx[0]], nums[idx[1]], n]))
                     for idx in two_sums[-n] if i not in idx}

    return sums

print(threeSumZero2([-1,0,1,2,-1,4]))

Output

{(-1, -1, 2), (-1, 0, 1)}

Upvotes: 3

Joe Iddon
Joe Iddon

Reputation: 20424

You could do this with itertools (see Oliver's answer), but you can also achieve the result with three nested for-loops:

def threeSumZero2(lst):
    groups = []
    for i in range(len(lst)-2):
        for j in range(i + 1, len(lst)-1):
            for k in range(j + 1, len(lst)):
                if lst[i] + lst[j] + lst[k] == 0:
                    groups.append((lst[i], lst[j], lst[k]))
    return groups

and your test:

>>> threeSumZero2([-1, 0, 1, 2, -1, 4])
[(-1, 0, 1), (-1, 2, -1), (0, 1, -1)]

Oh and list != array!

Upvotes: 0

Related Questions