CODEWITHSUNDEEP
pythonlistcombinationspython-itertools
DPdl
DPdl

Reputation: 755

Using Itertools to create a list of combination of elements from multiple lists

I have the following:

a = [i/100 for i in range(5,105,5)]
a.append(0)
b = [i/100 for i in range(5,105,5)]
b.append(0)
c = [i/100 for i in range(5,105,5)]
c.append(0)
d = [i/100 for i in range(5,105,5)]
d.append(0)
e = [i/100 for i in range(5,105,5)]
e.append(0)

combs =  itertools.product(a,b,c,d,e)

'combs' would give me all possible combinations of a,b,c,d and e. However, I was wondering if I could combine them so that they add up to 1.

Thank you.

Upvotes: 2

Views: 176

Answers (2)

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 477794

As far as I know there is no builtin way in to do this. You can evidently filter the result, but this would be rather inefficient: one expects only a small amount of combinations to add up to sum.

Given all fed values are positive (zero is acceptable) however, you can use this lazy function:

def product_sum(sum,*args):
    return product_sum_intern(sum,0,0,[],*args)

def product_sum_intern(sum,cur_sum,idx,cur,*args):
    if idx >= len(args):
        if sum == cur_sum:
            yield tuple(cur)
    elif cur_sum <= sum:
        for x in args[idx]:
            cur.append(x)
            for e in product_sum_intern(sum,cur_sum+x,idx+1,cur,*args):
                yield e
            cur.pop()

For instance:

>>> list(product_sum(15,[1,12],[1,4,7],[0,3,6,7],[0,1]))
[(1, 7, 6, 1), (1, 7, 7, 0)]

this algorithm gives up on branches as soon as it finds out it is already overshooting the sum. A more advanced algorithm exists that also gives up branches if there is no way to get in reach of the sum.

Upvotes: 1

Robᵩ
Robᵩ

Reputation: 168876

To produce a list of all such results:

combs = [comb for comb in combs if sum(comb) ==1]

Or, if you can use a generator of that data:

combs = (comb for comb in combs if sum(comb) ==1)

Upvotes: 0

Related Questions