Iterate over itertools.product in different order without ever creating list

Question

I have an iterable that covers a huge search space. My plan is not to let the script terminate, but to just kill it after a certain amount of time.

Now I need the Cartesian product of this space and search in there. itertools.product produces this order:

>>> list(itertools.product(range(3), repeat=2))
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]

While I want to search in a diagonal order similar to:

[(0, 0), (0, 1), (1, 0), (0, 2), (1, 1), (2, 0), (1, 2), (2, 1), (2, 2)]

sorted with some key function that returns the sum of the elements of the tuple would be my regular approach, but for sorting all data needs to be inspected, which is infeasible in my case. Is there a way to do this?

This question is very similar to this one, but there sorted is still used in the answer. Also I don't quickly see how to adapt ordered_combinations to ordered_product.

Problem statement is limited for `repeat=2`, because OP's thinking in terms of diagonals of square matrix M, where number of rows and columns is N (in this case 3). By accident `itertools.product` produces list of all element positions in matrix and OP's trying to manipulate those for expected output, but what really should be done here is completely different solution based on matrix-like problem statement. — Łukasz Rogalski, Oct 07 '16 at 13:00
@Chris_Rands I guess still in order of the sum of the elements, and lexicographically for an equal sum. Although I'm only interested in the repeat=2 case and I don't care about the direction of the diagonal, e.g., I don't care whether `(0,1)` or `(1,0)` comes first. @ŁukaszRogalski well that's how I pictured it, yes, it doesn't have to be the structure of the answer (although it could be). — qpllb, Oct 07 '16 at 13:02

score 1 · Answer 1 · answered Aug 06 '17 at 21:20

The question is equivalent to asking how to create all n-tuples for with a given sum for successive and increasing values of the sum:

                  (0, 0),               sum == 0
              (0, 1), (1, 0),           sum == 1
        (0, 2), (1, 1), (2, 0),         sum == 2
             (1, 2), (2, 1),            sum == 3
                  (2, 2)                sum == 4

For any given row (with a given target sum), the subproblem is equivalent to the dynamic programming problem Number of ways to make change for amount N or Number of ways to add up to a sum S with N numbers .

Also see the write ups in Combinatorial Algorithms by Donald Knuth.

score 0 · Answer 2 · answered Aug 06 '17 at 21:54

It's actually quite easy to calculate the indices for the diagonals for the repeat=2 case:

def diagonal_product(inp):
    inp = tuple(inp)
    n = len(inp)
    # upper left triangle
    for i in range(n):
        for j in range(i+1):
            yield inp[i-j], inp[j]
    # lower right triangle
    for i in range(1, n):
        for j in range(n-i):
            yield inp[n-j-1], inp[i+j]

Just to show the results:

>>> list(diagonal_product(range(4)))
[(0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (0, 2), (3, 0), (2, 1), (1, 2), 
 (0, 3), (3, 1), (2, 2), (1, 3), (3, 2), (2, 3), (3, 3)]

>>> list(diagonal_product(range(3)))
[(0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (0, 2), (2, 1), (1, 2), (2, 2)]

The function itself might be a bit more complicated (or slower) than it needs to be but I couldn't find any reference implementation for this case just yet.

In case of a range as input you could also avoid all the indexing and just return the indices:

def diagonal_product(n):
    # upper left triangle
    for i in range(n):
        for j in range(i+1):
            yield i-j, j
    # lower right triangle
    for i in range(1, n):
        for j in range(n-i):
            yield n-j-1, i+j

list(diagonal_product(3))
# [(0, 0), (1, 0), (0, 1), (2, 0), (1, 1), (0, 2), (2, 1), (1, 2), (2, 2)]

Iterate over itertools.product in different order without ever creating list

2 Answers2