Python list comprehension two for loops

When working with lists in Python, you'll likely often find yourself in situations where you'll need to translate values from one list to another based on specific criteria.

Nội dung chính Show

Quick Refresher: What Is a List Comprehension?
What Is a Nested List Comprehension?
Before You Start Nesting...
Example List comprehension nested for loop
Nested List Comprehension to flatten a given 2-D matrix

Generally, if you're working with small datasets, then using for loops instead of list comprehensions isn't the end of the world. However, as the sizes of your datasets start to increase, you'll notice that working through lists one item at a time can take a long time.

Let's generate a list of ten thousand random numbers, ranging in value from one to a million, and store this as num_list. We can then use a for loop and a list comprehension to generate a new list containing the num_list values greater than half a million. Finally, using %timeit, we can compare the speed of the two approaches:

When you’re learning of a new programming language, it’s challenging to figure out how to keep your code clear.

One of the best ways to make Python code more accessible to your colleagues comes from collapsing multiple lines of code into a single line. In a recent article I showed you how to do this using list comprehensions. This article builds off that one, so you should definitely check it out before continuing here.

As it turns out, you can nest list comprehensions within another list comprehension to further reduce your code and make it easier to read still. As a matter of fact, there’s no limit to the number of comprehensions you can nest within each other, which makes it possible to write very complex code in a single line.

Since this is such a powerful function, let’s take a deeper look.

List comprehensions create a new list by scanning all items in a list, checking to see if they meet a condition, passing the values that do and applying a function to them. You can place one list comprehensions within another list comprehension. This creates a nested list comprehension and further simplifies your code.

Back to BasicsIncrease the Readability of Your Python Script With 1 Simple Tool

Quick Refresher: What Is a List Comprehension?

As a quick refresher, a list comprehension is a way to create a new list by filtering elements in a list that meet a condition, transforming the elements as desired and storing the new values in a list. All we need is a single line of code. The code is as follows:

Result = [Function for Value in DataSet if Condition]

That code sorts through each value in your data set, checks to see if it meets the condition, applies the stated function and stores it in a list named Result. It applies a for loop, if statement and a function all in a single line.

Here’s an example showing how list comprehensions work:

Numbers = [1,2,3,4,5,6,7,8,9,10] Result_ForLoop = [] for Number in Numbers: if Number > sum(Numbers)/len(Numbers): Result.append(Number) Result_ListComprehension = [] Result_ListComprehension = [Number for Number in Numbers if Number > sum(Numbers)/len(Numbers)] print(Result_ForLoop) print(Result_ListComprehension)

Those two print statements will show the exact same result: [6,7,8,9,10].

For a more detailed discussion of how those two pieces of code return the same result see my guide to list comprehensions.

Now let’s dive in to nested list comprehensions.

What Is a Nested List Comprehension?

A nested list comprehension doubles down on the concept of list comprehensions. It’s a way to combine not only one, but multiple for loops, if statements and functions into a single line of code. This becomes useful when you have a list of lists (instead of merely a single list).

Let’s consider an example. Say you have two lists of numbers and you want to return the square of all of the even numbers. You could write the following code using a single list comprehension:

Numbers = [[1,2,3,4,5,6,7,8,9,10], [11,12,13,14,15,16,17,18,19,20]] Result = [] for list in Numbers: Squares = [Number ** 2 for Number in list if Number % 2 == 0] Result.extend(Squares)

The above code:

Creates the input data set and stores it in Numbers. The input data consists of two lists. The first list runs from one to 10, and the second list runs from 11 to 20.
Then the code creates an empty list called Result. We will use Result to store the output from our function.
The code will then create a for loop and iterate through each list in Number.
Within the for loop it uses a list comprehension to search through the provided list, check if each number in the list is divisible by two (i.e. Is it even?), squares the result and stores it in a list. This list is then stored in the variable Squares.
The final line then adds Squares to Result giving us a list of the squares of the values that are divisible by two.

Since there are two lists in Numbers it executes the for loop twice.

Now, since you’re familiar with list comprehensions you probably understand it’s fully possible to remove the for loop and if statement with a list comprehension. Doing so creates a nested list comprehension.

To do this we write a single list comprehension with two for loops. It’s key to remember three things when doing this:

The function is always the first term.
The for loops are always written in the order of the nesting.
The condition is always placed at the end.

So, to construct a nested for loop in our example we need to:

Write the function to square the numbers that have passed the filter.
Write a for loop iterating through all of the lists in Numbers.
Write a for loop iterating through each number in the passed list.
Write a condition passing the numbers that are even.

And we need to do it in that order. Fortunately, once we’ve created that structure in our minds it’s easy to do. The code appears as follows:

Result = [Number ** 2 for list in Numbers for Number in list if Number %2 == 0]

If you replace the previous code you’ll find this nested list comprehension returns the same result.

This is a powerful way to reduce several lines of code into a single line and can significantly increase the readability of your code. Since it’s in a single line, anybody who understands list comprehensions will be able to quickly deduce what is happening and follow your logic.

Learn More With Peter GrantNeed to Automate Your Data Analysis? Here’s How.

Before You Start Nesting...

While nested list comprehensions are a useful way to improve your code's readability, there is a downside. List comprehensions can get very complex very quickly. What if you have a hundred for loops? Or what if you have very complex function and condition descriptions?

I could envision a complex nested for loop spanning several lines of code, and becoming nearly impossible to read or understand. While there’s Python has no technical limit to understanding complex a nested for loops, there is a human limit.

If you have a complex piece of code with many for loops, or complex functions and conditions, using a nested for loop may actually make your code more difficult to understand. Take care to ensure your efforts are making your collaborators’ lives easier instead of harder.

Generally speaking, if you have more than three levels of nesting it may be easier for everybody if you just write out the for loops.

Learning to make code easy for other people to understand is a critical task in the modern workplace —especially as more and more teams work remotely. List comprehensions and nested list comprehensions are a useful Python tool to help your collaborators quickly understand how your code works.