How to generate itertools.product with a “sum” condition?

My code is:

        import itertools
        a = [*range(22)], [*range(22)], [*range(22)], [*range(22)]

        combination = [seq for seq in itertools.product(*a) if sum(seq) <= 21]
        print(combination)

        lst1 = [item[0] for item in combination]
        lst2 = [item[1] for item in combination]
        lst3 = [item[2] for item in combination]
        lst4 = [item[3] for item in combination]

My result is:

“[….., (0, 0, 0, 19), (0, 0, 0, 20), (0, 0, 0, 21), (0, 0, 1, 0), (0, 0, 1, 1), (0, 0, 1, 2), (0, 0, 1, 3), (0, 0, 1, 4), ….., (0, 0, 1, 16), (0, 0, 1, 17), (0, 0, 1, 18), (0, 0, 1, 19), (0, 0, 1, 20),….]”

The cap in the generated set of iterables is that the elements in each iterable should be less than or equal to 21.

But what I want is the cap to be on the first two and last two elements of each iterable.

For example, if the current cap results in (0, 0, 1, 20) because it sums to 21; I want the cap to instead be on the first two and last two elements summing to 21 separately resulting in (0, 21, 1, 20).

How do I do that?

Answer

Deconstruct the problem to only the clause in question. This has nothing to do with itertools.

if sum(seq) <= 21

Simply write the condition you described:

if sum(seq[:2] ) <= 21 and
   sum(seq[-2:]) <= 21

Advertisement

Answer