The world’s leading publication for data science, AI, and ML professionals.

All About Python List Comprehension

Elegant, comfortable, concise, and fast way to build lists

Image Courtesy: Pixabay
Image Courtesy: Pixabay

Python is a language of expressiveness in concise code, which is very elegant and easy to comprehend. List Comprehension provides a concise way to create lists. List Comprehension is the proper "pythonic" way of building an accessible, concise, and fast way to build lists. It allows one to perform complex operations on lists using a single line.

What are List Comprehensions?

List comprehensions provide us with a simple way to create a list based on some sequence or another list that we can loop over. In Python terminology, anything that we can loop over is called iterable. At its most basic level, list comprehension is a syntactic construct for creating lists from existing lists. In the core of list comprehension, it is the loop (‘for’ loop). Any list comprehension we can represent as a for loop, but when we represent it with equivalent list comprehension in a single line, it looks genuinely unique.

List comprehension has these essential parts :

  1. An iterable input sequence (this could be a list, a range, or any sequence) that we iterate using a variable name.
  2. An output expression
  3. An optional condition for the variable to filter or map or do some logical action

For example, if we want to create a square list from a list of variables from 1 to 5, we can do so as below using the standard for loop, but when we use a list comprehension, it becomes concise and lightweight.

Image: Weighing machine from Pixabay
Image: Weighing machine from Pixabay

If we break-down the list comprehension parts, it looks as below.

Explanation of a simple list comprehension that output square of each input variable
Explanation of a simple list comprehension that output square of each input variable

Thus here iterable is range(1,6), and x is the variable representing each value in the range. For each x (1,2,3,4,5) , output expression x*x calculates a corresponding output (1,4,9,16,25) which ultimately becomes the final output list (refer to above diagram).

Few Ways of Usage other than the basic construct:

1) Filter based on condition

List comprehensions provide a way to check for conditions and filter out the elements we don’t need from the input list for output. This filtering happens before the output expression takes place. Refer to the below example where we calculate square for only those variables which are divisible by 2.

List Comprehension with Filtering Condition
List Comprehension with Filtering Condition

2) Nested Condition

List Comprehension can be used to check nested conditions. Refer to the below code, where we want to find which all variables in the input list are divisible by 2,4 & 5.

3) if..else in List Comprehension

We can use an "if-else" in a list comprehension in Python. Refer to the code block below, which finds odd and even for a list of numbers. We need to specify what the output is first, following "if" condition and then "else" followed by what is the output under "else."

Note, we can not print input value directly without using a comma and parenthesis.

4) Flattening a multi-dimensional list

Suppose we have to flatten a 2D list. We can easily do it using List Comprehension using a sublist.

How can we flatten a 3D list? It becomes a bit complex where first we flatten 3D to 2D and then 2D to 1D. In the below example blue underlined portion is the step that makes 2D flatten from a 3D list, and then the outer part does flattening to 1D.

5) Nested "for" List Comprehension

We can use list comprehension for multiple nested "for" loops. We have to careful in terms of the order of execution when we use nested list comprehension. Suppose we have a 3×3 square matrix, and we want to create a matrix using this matrix where all diagonal values need to be square while other values need to be zeros. Here the output is 1 dimensional, and we need additional logic to convert to 2D matrix form. Here I used "column" and "row" for explanation purposes but can be used directly in the "for loop" without the intermediate variables (as used in below List Comprehension).

matrix = [[ 1, 2, 3],
          [ 4, 5, 6],
          [ 7, 8, 9]]
row = len(matrix)
column = len(matrix[0])
output = []
for i in range(row):
    for j in range(column):
        if i == j:
            output.append(matrix[i][j] * matrix[i][j])
        else:
            output.append(0)

output
[1, 0, 0, 0, 25, 0, 0, 0, 81]

Now we see equivalent nested list comprehension (single line only, but for easy understanding, I have made it in 3 lines).

[[matrix[i][j]*matrix[i][j] if i == j else 0 
              for i in range(len(matrix))] 
                  for j in range(len(matrix[0]))]

Now let’s understand this with a picture of how each loop is working. Loop execution order is top to bottom (bold green arrow). Inner loops (dark green ) we complete first, and then control comes back to the outer loop (blue). Here List Comprehension is one line in comparison to 5 lines of ordinary "for loop," as shown above code. We can also get the matrix form of output directly by putting brackets after first for loop.

Advantages of List Comprehension

  1. One of the significant advantages of list comprehension is that developer need to write less code
  2. List comprehension often becomes easier to understand
  3. The speed of list comprehensions is notably better than for-loops when appending items to the list
  4. List comprehensions are an excellent alternative to the built-in map and filter functions

Other Comprehensions

In Python, we have also "dictionary comprehensions" and "set comprehensions" that are similar in operation but meant for dictionary and set. Python also has "generator expression" which uses a "round" bracket.

Final Note

When I started my journey with Python, I used to find it challenging to understand how the list comprehension works. Once I understood its elegance and power, I never looked back. Even though it feels intimidating at first, with a bit of practice, every python lover can fall in love with List Comprehension.

I look forward to your comment and share if you have any unique experience with the list comprehension. Thanks for reading. You can connect me @LinkedIn.


For only $5/month, get unlimited access to the most inspiring and uplifting content… Click on the link below to become a Medium member and support my writing. Thank you! https://baijayanta.medium.com/membership


Related Articles