user3658033 user3658033 - 1 month ago 9
Python Question

Itertools.product to customize the amount of combos for each input?

I'm using itertools.product to come up with a combination of groups.I'm bad at explaining without examples, so here's code as is.

group1=[1,2,3];group2=[4,5,6];group3=[7,8,9]
list(itertools.product(group1,group2,group3))


This gives me all combinations of 1 from each group. But how would I go about getting combinations of 2 numbers from group 1, 2 numbers from group 2, and 1 number from group 3?

So for example, I'd like the combination (1,2,5,6,9) to be in the list. Is it possible to customize this? itertools.product doesn't seem to be as flexible as I need it to be, and I've been unsuccessful in learning cartesian products enough to understand how to tweak the .product function.

EDIT: I made the groups small to keep it simple, but each group will have hundreds of unique values.

Answer

Take the cartesian product of the r-combinations of each group:

from itertools import product, chain, combinations, permutations

groups = [[1,2,3],[4,5,6],[7,8,9]]
counts = (2, 2, 1)

selections = [combinations(g, c) for g, c in zip(groups, counts)]

for n_tuple in product(*selections):
    print(tuple(chain.from_iterable(n_tuple)))

Output:

(1, 2, 4, 5, 7)
(1, 2, 4, 5, 8)
(1, 2, 4, 5, 9)
(1, 2, 4, 6, 7)
(1, 2, 4, 6, 8)
(1, 2, 4, 6, 9)
(1, 2, 5, 6, 7)
(1, 2, 5, 6, 8)
(1, 2, 5, 6, 9)
(1, 3, 4, 5, 7)
(1, 3, 4, 5, 8)
(1, 3, 4, 5, 9)
(1, 3, 4, 6, 7)
(1, 3, 4, 6, 8)
(1, 3, 4, 6, 9)
(1, 3, 5, 6, 7)
(1, 3, 5, 6, 8)
(1, 3, 5, 6, 9)
(2, 3, 4, 5, 7)
(2, 3, 4, 5, 8)
(2, 3, 4, 5, 9)
(2, 3, 4, 6, 7)
(2, 3, 4, 6, 8)
(2, 3, 4, 6, 9)
(2, 3, 5, 6, 7)
(2, 3, 5, 6, 8)
(2, 3, 5, 6, 9)

You can change combinations to permutations if order matters when selecting from each group (for example, if (3, 2, 5, 6, 9) is different from (2, 3, 5, 6, 9)).

You should note that this produces choose(|g1|, c1) * choose(|g2|, c2) * ... * choose(|gN|, cN) elements from N groups, where choose(n, k) is the binomial coefficient. This is incalculably large if your group sizes are in the hundreds as you say--or if the number of groups is large, as well.