CODEWITHSUNDEEP
pythonrecursion
Luke Gregor
Luke Gregor

Reputation: 133

Variable for loops with recursion

Would like to do the following by recursion so that I can vary the number of 'for' loops:

n = 5
out = []
for i in range(n):
    for j in range(i,n):
        for k in range(j,n):
            out.append([i,j,k])

To return

out =   [[0 0 0]
         [0 0 1]
         [0 0 2]
         [0 0 3]
         [0 0 4]
         [0 1 1]
         [0 1 2]
         [0 1 3]
         [0 1 4]
         [0 2 2]
         [0 2 3]
         [0 2 4]
         [0 3 3]
         [0 3 4]
         [0 4 4]
         [1 1 1]
         [1 1 2]
         [1 1 3]
         [1 1 4]
         [1 2 2]...]

e.g.

def Recurse(n, p):
  # where p is the number of for loops
  some magic recursion 
  return out

I've had a look at some of the other recursion questions, but struggling to get to the solution.

Upvotes: 2

Views: 175

Answers (2)

John La Rooy
John La Rooy

Reputation: 304205

@DSM has a better answer (but in the comments)

However here is a simple recursive solution, in case anyone is struggling to work it out

def f(n, p, start=0):
    if p==0:
        yield []
    else:
        for i in range(start, n):
            for j in f(n, p-1, i):
                yield [i]+j

Upvotes: 2

Andrew Clark
Andrew Clark

Reputation: 208475

Instead of using recursion, use itertools.product(), which is equivalent to nested for loops in a generator expression:

>>> import itertools
>>> 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)]
>>> list(itertools.product(range(3), repeat=3))
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 2, 0), (0, 2, 1), (0, 2, 2), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 2, 0), (1, 2, 1), (1, 2, 2), (2, 0, 0), (2, 0, 1), (2, 0, 2), (2, 1, 0), (2, 1, 1), (2, 1, 2), (2, 2, 0), (2, 2, 1), (2, 2, 2)]

edit: Didn't realize that this isn't actually a Cartesian product since the inner loops use the outer variable to start the range, here is one possibility but it is not as efficient as it could be because it generates extra values and needs to check that each value is valid:

def nested_loops(n, num_loops):
    prod = itertools.product(range(n), repeat=num_loops)
    for item in prod:
        if all(item[i] <= item[i+1] for i in range(num_loops-1)):
            yield item

>>> list(nested_loops(3, 2))
[(0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2)]
>>> list(nested_loops(3, 3))
[(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 1), (0, 1, 2), (0, 2, 2), (1, 1, 1), (1, 1, 2), (1, 2, 2), (2, 2, 2)]

Upvotes: 4

Related Questions