CODEWITHSUNDEEP
pythonalgorithmloopsnested-loops
dhack
dhack

Reputation: 23

Next level of abstraction up from a for loop?

Working in Python, but looking for a general structure.

I have a function that takes as input an array of length L (e.g., [0,0,0,...,0] with L semi-positive elements). Each element of the input array has a positive maximum value, e.g., I have an array of length L, max_vals = [10,5,4,...,7]. I would like to iterate over all possible inputs to the function.

This can be done pretty easily using "generalized counting", but I wonder if there's a more graceful way to do this. Specifically, I'd like some structure like:

meta_for vector in meta_range(max_vals):
    relevant_function( vector )

That is functionally equivalent to:

for i_1 in range(max_vals[0]):
    for i_2 in range(max_vals[1]):
        ...
            for i_L in range(max_vals[L-1]):
                relevant_function( [i_1, i_2, ..., i_L] )

The difficulty is that the number of nested for loops depends on L, and so can't be hardcoded.

Does anything like this exist?

Upvotes: 2

Views: 237

Answers (1)

Peter de Rivaz
Peter de Rivaz

Reputation: 33509

You can do this with itertools.product:

import itertools
max_vals = [3, 2]
for I in itertools.product(*[range(m) for m in max_vals]):
    print I

Prints

(0, 0)
(0, 1)
(1, 0)
(1, 1)
(2, 0)
(2, 1)

Upvotes: 3

Related Questions