Search code examples
pythonalgorithmloopsnested-loops

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?


Solution

  • 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)