[trick]使用生成器打破嵌套循环

news2024/11/25 18:40:45

原文

chain

break用于结束循环。但是，如果有嵌套循环，如何跳出外层循环？

def this_is_the_one(x):
    return x == 3

my_list = [[1, 2], [3, 4], [5, 6]]
for sublist in my_list:
    for element in sublist:
        print(f"Checking {element}")
        if this_is_the_one(element):
            # ...?

break 不会有太大帮助，因为它只会跳出内部 for 循环：

for sublist in my_list:
    for element in sublist:
        print(f"Checking {element}")
        if this_is_the_one(element):
            break

"""Output:
Checking 1
Checking 2
Checking 3  # 只跳过了 4
Checking 5
Checking 6
"""

有多种方法可以解决这个困难，在本短文中，我想向您展示如何使用生成器来解决这个问题。核心思想是将嵌套循环提取到一个单独的生成器中，然后在原始循环的位置使用该生成器。
将嵌套循环变为一个生成器：

def elements_from_sublists(my_list):
    for sublist in my_list:
        for element in sublist:
            yield element

使用生成器：

for element in elements_from_sublists(my_list):
    print(f"Checking {element}")
    if this_is_the_one(element):
        break

"""Output:
Checking 1
Checking 2
Checking 3
"""

现在我们只有一个循环，所以关键字 break 运行得很完美。将循环逻辑提取到生成器中的这种策略可能非常有用，所以要留意使用它的机会！

奖励：
已经有标准库函数实现了上述的elements_from_sublists功能：

from itertools import chain

for element in chain.from_iterable(my_list):
    print(f"Checking {element}")
    if this_is_the_one(element):
        break