[Memo] Several Ideas for Breaking Out of Nested Loops in Python

【备忘】Python中断多重循环的几种思路 苏剑林 · 2016-12-19 · Machine Learning Python memo

【Note to self】Several ways to break out of nested loops in Python

Breaking out of a single loop

Whatever programming language you're using, you'll sometimes want to break out of a loop — for example, when enumerating values, you might want to stop as soon as you find one that satisfies some condition. Breaking out of a single loop is easy, e.g.

for i in range(10):
    if i > 5:
        print i
        break

Sometimes, though, we need to break out of nested loops, and break only exits one level, e.g.

for i in range(10):
    for j in range(10):
        if i+j > 5:
            print i,j
            break

Code like this doesn't stop as soon as it finds one pair with i+j > 5 — instead it keeps finding 10 such pairs in a row, because break only breaks out of the for j in range(10) loop. So how do we break out of multiple levels? Let me jot down some notes here. more

Breaking out of nested loops

As it happens, standard Python syntax doesn't support breaking out of nested loops directly, so we have to resort to some tricks. The main approaches are: wrapping the loops in a function, using a Cartesian product, or exploiting debugging mode.

Wrapping it in a function

In Python, a function stops running as soon as it hits a return statement. We can exploit this to terminate nested loops by wrapping the logic in a function, for example

def work():
    for i in range(10):
        for j in range(10):
            if i+j > 5:
                return i,j

print work()

Using a Cartesian product

The idea here is: since we can break out of a single loop, why not rewrite the nested loops as a single loop? We can do this with the Cartesian product function product from itertools, for example

from itertools import product

for i,j in product(range(10), range(10)):
    if i+j > 5:
        print i,j
        break

Using debugging mode

The Cartesian product trick is clever and concise, but it only works when the sets being looped over at each level are independent of each other. If each level of the loop depends closely on the previous one, this trick won't work. In that case, you can go back to the first approach and wrap it in a function — or, alternatively, exploit debugging mode. This approach relies on the fact that debugging mode exits as soon as an error occurs, so it deliberately triggers a fake error.

class Found(Exception):
    pass

try:
    for i in range(10):
        for j in range(i): #第二重循环跟第一重有关
            if i + j > 5:
                raise Found
except Found:
    print i, j