Python List Comprehension
A quick and easy introduction to Python list comprehension
Introduction
In this article, I want to show you a very useful feature of the Python language: list comprehension. After reading, you will be able to write your code more efficiently and beautifully.
List comprehension is an elegant way to define and create lists based on existing lists or iterative objects.
Examples
Basic usage
Let’s imagine that we need to create the list of squared numbers from 0 to 10. I think if you are not familiar with list comprehension, you will do it like in the code below.
But wait, we can use the power of Python language, and do it more elegant:
If you check the arr values at the end of the programs you will see the same values. But with list comprehension, the code looks more compact and it’s easy to read.
Structure
The main structure of this operation have the next form:
[some_processed_variable for some_variable in iterable_object]If condition
Let’s go deeper and add additional functionality to this structure. Suppose we need to get a list of only odd numbers. Now we can do it like in the code below.
We added an additional “if” condition to the end of the structure (of course we can do it by setting range argument, like range(1, 10, 2)).
Nested loops
We can write a nested loop in the list comprehension. For this, we write two loops one by one — second will be nested. Let’s check it with the common example — create a card deck where each card has suit and value.
As you can see we have two loops: main which iterate over suits and nested which iterate over a list of values.
We can write more nested loops, but if it becomes difficult to read and understand it is best not to.
Another good example of a nested loop with comprehension is creating a multiplication table. Below we create a list with strings in the format “number1 x number2= number1 * number2”. Where numbers 1 and 2 are numbers from 1 to 9.
Real example
The final example I met in one of the Kaggle competitions. There is a table with a lot of columns. And some of them have the format “x<some_number>”. And I wanted to get all these column names for future processing. With list comprehension, it was so easy. There is a small example of this task and solution for it in the next code cell.
Conclusions
In this article, we focused on the list comprehension mechanism in Python. The main concept is it is easy for usage and understanding. But we need to be careful with adding more nested loops and difficult conditionals to our structure — it can start to be messy.
Besides list comprehension, Python has dictionary comprehension and set comprehension. The main difference between them and the list comprehension is to use {} brackets instead of []. And for the dictionary, we need to specify the key with the value.
