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

Question

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 1

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

Answer 2

It looks to me like what you really need is this:

import itertools

combinations = [
    seq
    for seq in itertools.product(*itertools.tee(range(22), 4))
    if sum(seq[:2]) == 21 and sum(seq[-2:]) == 21
]
print(combinations)

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

As you said that you wanted the cap to be when the sum of the first 2 or last 2 elements of the seq was equal to 21, not lower or equal.