Python Language List Comprehensions List Comprehensions with Nested Loops
Example
List Comprehensions can use nested for loops. You can code any number of nested for loops within a list comprehension, and each for loop may have an optional associated if test. When doing so, the order of the for constructs is the same order as when writing a series of nested for statements. The general structure of list comprehensions looks like this:
[ expression for target1 in iterable1 [if condition1]
for target2 in iterable2 [if condition2]...
for targetN in iterableN [if conditionN] ]
For example, the following code flattening a list of lists using multiple for statements:
data = [[1, 2], [3, 4], [5, 6]]
output = []
for each_list in data:
for element in each_list:
output.append(element)
print(output)
# Out: [1, 2, 3, 4, 5, 6]
can be equivalently written as a list comprehension with multiple for constructs:
data = [[1, 2], [3, 4], [5, 6]]
output = [element for each_list in data for element in each_list]
print(output)
# Out: [1, 2, 3, 4, 5, 6]
In both the expanded form and the list comprehension, the outer loop (first for statement) comes first.
In addition to being more compact, the nested comprehension is also significantly faster.
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
The overhead for the function call above is about 140ns.
Inline ifs are nested similarly, and may occur in any position after the first 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)
# Out: [2, 3, 4]
For the sake of readability, however, you should consider using traditional for-loops. This is especially true when nesting is more than 2 levels deep, and/or the logic of the comprehension is too complex. multiple nested loop list comprehension could be error prone or it gives unexpected result.
