Python嵌套循环的列表理解

示例

列表推导可以使用嵌套for循环。您可以在列表推导中编码任意数量的嵌套for循环,并且每个for循环都可以具有可选的关联if测试。这样做时,for构造的顺序与编写一系列嵌套for语句时的顺序相同。列表推导的一般结构如下:

[ expression for target1 in iterable1 [if condition1]
             for target2 in iterable2 [if condition2]...
             for targetN in iterableN [if conditionN] ]

例如,以下代码使用多个for语句将列表的列表展平:

data = [[1, 2], [3, 4], [5, 6]]
output = []
for each_list in data:
    for element in each_list:
        output.append(element)
print(output)
# 输出:[1、2、3、4、5、6]

可以等效地写成具有多种for构造的列表理解:

data = [[1, 2], [3, 4], [5, 6]]
output = [element for each_list in data for element in each_list]
print(output)
# 输出:[1、2、3、4、5、6]

现场演示

在扩展形式和列表理解中,外部循环(对于语句而言是第一个)都是第一位的。


除了更紧凑之外,嵌套理解还明显更快。

In [1]: data = [[1,2],[3,4],[5,6]]
In [2]: def f():
   ...:     output=[]
   ...:     for each_list in data:
   ...:         for element in each_list:
   ...:             output.append(element)
   ...:     return output
In [3]: timeit f()
1000000 loops, best of 3: 1.37 µs per loop
In [4]: timeit [inner for outer in data for inner in outer]
1000000 loops, best of 3: 632 ns per loop

以上函数调用的开销约为140ns


内联ifs类似地嵌套,并且可以在第一个之后的任何位置出现for:

data = [[1], [2, 3], [4, 5]]
output = [element for each_list in data
                if len(each_list) == 2
                for element in each_list
                if element != 5]
print(output)
# 出:[2、3、4]

现场演示

但是,出于可读性考虑,应考虑使用传统的for循环。当嵌套深度超过2级时,和/或理解的逻辑过于复杂时,尤其如此。多个嵌套循环列表理解可能容易出错,或者产生意外结果。