CODEWITHSUNDEEP
pythonpython-3.xpython-itertoolsfunctools
qedpi
qedpi

Reputation: 43

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.

Upvotes: 4

Views: 355

Answers (2)

Stefan Pochmann
Stefan Pochmann

Reputation: 28606

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

Upvotes: 3

Laurent LAPORTE
Laurent LAPORTE

Reputation: 22952

Here is an example of what you want:

B = [10, 15, 20, 5]

F = [lambda x: x,
     lambda x: x * x,
     lambda x: x * 2 - 5]

def dostuff(i0, i1, i2, i3):
    print((i0, i1, i2, i3))

expected = []
for i0 in range(0, B[0]):
    for i1 in range(F[0](i0), B[1]):
        for i2 in range(F[1](i1), B[2]):
            for i3 in range(F[2](i2), B[3]):
                expected.append([i0, i1, i2, i3])

I found a recursive solution like this:

def iter_rec(found, fL, bL):
    if fL and bL:
        ik = found[-1] if found else 0
        fk = fL[0]
        bk = bL[0]
        for i in range(fk(ik), bk):
            for item in iter_rec(found + [i], fL[1:], bL[1:]):
                yield item
    else:
        yield found

# prepend the null function to ensure F and B have the same size
F = [lambda x: 0] + F

current = [item for item in iter_rec([], F, B)]

We have the same result.

assert expected == current

Upvotes: 2

Related Questions