CODEWITHSUNDEEP
pythonoptimizationfor-loop
user3024444
user3024444

Reputation: 13

What is the best way to write a multiple nested for loops in Python

I am working on a code where I need to use a nested loop as follows:

for r in range(m):
   for s in range(r+1, m):
      for t in range(s+1, m):
         for u in range(t+1, m):
            for v in range(u+1, m):
                  arr.append([r,s,t,u,v])

But this traditional nested loops look pretty ugly. Is there a way to perform the same operation with less lines?

I looked at itertools.product but I was not able to get what I want since all start/end indices of my loops depends on the previous level.

Upvotes: 1

Views: 1462

Answers (3)

thefourtheye
thefourtheye

Reputation: 239453

You can use itertools.combinations and the second parameter is the number of loops you wanted to execute.

from itertools import combinations
for item in combinations("ABCD", 3):
    print item

Output

('A', 'B', 'C')
('A', 'B', 'D')
('A', 'C', 'D')
('B', 'C', 'D')

So, with list comprehension, the whole code becomes like this

[list(item) for item in combinations("ABCD", 3)]

Upvotes: 5

6502
6502

Reputation: 114481

Sometimes this kind of algorithm in general can be implemented using a recursive approach. As an example where the iteration for next level requires complex testing consider an eight-queens problem solver:

for col0 in range(8):
    for col1 in range(8):
        if col1 != col0 and not same_diag(0, col0, 1, col1):
            for col2 in range(8):
                if (col2 != col1 and
                    col2 != col0 and
                    not same_diag(0, col0, 2, col0) and
                    not same_diag(1, col1, 2, col2)):
                    for col3 in range(8):
                        ... same pattern up to col7 ...

this repetitive code can be factored out with

def find_solution(i, free_cols, placed_queens):
    if i == 8:
        print_solution(placed_queens)
    else:
        for col in free_cols:
            if not any(same_diag(i, col, j, q)
                       for j, q in enumerate(placed_queens)):
                find_solution(i+1,
                              free_cols - set([col]),
                              placed_queens + [col])

and you also get that the number of levels (i.e. the board size for the queens problem) becomes a parameter.

Upvotes: 0

Steve Barnes
Steve Barnes

Reputation: 28370

Create your array directly from itertools.combinations, e.g. using n = 8:

>>> from itertools import combinations
>>> arr = [i for i in combinations(range(8), 5)]
>>> arr
[(0, 1, 2, 3, 4), (0, 1, 2, 3, 5), (0, 1, 2, 3, 6), (0, 1, 2, 3, 7), (0, 1, 2, 4, 5), (0, 1, 2, 4, 6), (0, 1, 2, 4, 7), (0, 1, 2, 5, 6), (0, 1, 2, 5, 7), (0, 1, 2, 6, 7), (0, 1, 3, 4, 5), (0, 1, 3, 4, 6), (0, 1, 3, 4, 7), (0, 1, 3, 5, 6), (0, 1, 3, 5, 7), (0, 1, 3, 6, 7), (0, 1, 4, 5, 6), (0, 1, 4, 5, 7), (0, 1, 4, 6, 7), (0, 1, 5, 6, 7), (0, 2, 3, 4, 5), (0, 2, 3, 4, 6), (0, 2, 3, 4, 7), (0, 2, 3, 5, 6), (0, 2, 3, 5, 7), (0, 2, 3, 6, 7), (0, 2, 4, 5, 6), (0, 2, 4, 5, 7), (0, 2, 4, 6, 7), (0, 2, 5, 6, 7), (0, 3, 4, 5, 6), (0, 3, 4, 5, 7), (0, 3, 4, 6, 7), (0, 3, 5, 6, 7), (0, 4, 5, 6, 7), (1, 2, 3, 4, 5), (1, 2, 3, 4, 6), (1, 2, 3, 4, 7), (1, 2, 3, 5, 6), (1, 2, 3, 5, 7), (1, 2, 3, 6, 7), (1, 2, 4, 5, 6), (1, 2, 4, 5, 7), (1, 2, 4, 6, 7), (1, 2, 5, 6, 7), (1, 3, 4, 5, 6), (1, 3, 4, 5, 7), (1, 3, 4, 6, 7), (1, 3, 5, 6, 7), (1, 4, 5, 6, 7), (2, 3, 4, 5, 6), (2, 3, 4, 5, 7), (2, 3, 4, 6, 7), (2, 3, 5, 6, 7), (2, 4, 5, 6, 7), (3, 4, 5, 6, 7)]
>>>

If you are looking for large numbers you might be better of writing a generator to iterate it for you or using the provided generator directly:

arr = [i for i in combinations(range(8), 5)]
for i in arr:
   # whatever you were going to use arr for....

Upvotes: 0

Related Questions