Search code examples
pythonpython-3.xpython-itertoolsfunctools

Using itertools for arbitrary number of nested loops of different ranges with dependencies?


Given a list of upperbounds: B1, B2, .. BN;
Dependency Functions: f1, ..., fN-1,

I'm wondering if there's a recipe using itertools or other classes in python for:

for i1 in range(0, B1):  
    for i2 in range(f1(i1), B2): ...
         for iN in range(fN-1(iN-1), BN)
             dostuff(i1, i2, ... iN)

Where there are N levels of nesting?
I want to use this helper function like this:
dependentProducts(Bs, fs, dostuff),
which returns a list or iterable

Ideally, the implementation would be iterative instead of recursive.


Solution

  • An iterative solution using @LaurentLAPORTE's setup. Put this code right under his and it should work. My args is a stack of the arguments fed into dostuff whenever it's full. The actual solution is the middle part, top and bottom parts are just testing.

    stefan = []
    def dostuff(*args):
        stefan.append(list(args))
    
    args = [-1]
    while args:
        n = len(args)
        args[-1] += 1
        if args[-1] >= B[n-1]:
            args.pop()
        elif n == len(B):
            dostuff(*args)
        else:
            args.append(F[n](args[-1]) - 1)
    
    assert expected == stefan