List comprehension is a concise way of writing a for loop on a list or a string. Performing an action of each element of a list or string and generation of a new list is very easy, short, and compact using list comprehension. Normally we use a for loop for that. Doing a list comprehension is a comparatively faster and shorter way of writing a ‘for loop’. I am assuming you already know python programming at list beginner level.

In this article, I will explain list comprehension with a lot of examples starting from the simplest example and keep moving towards the more complex problems.

Let’s start by creating a simple list of digits:

digits = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Let’s make a new list that includes the square of each element of the digits list. I will do it using a regular for loop first and then using a list comprehension.

Here is the regular for loop:

Example 1:

numbers = []for i in digits:
    numbers.append(i**2)

print(numbers)

Output:

[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Now the same thing using a list comprehension:

[i**2 for i in digits]

Output:

[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

As you can see instead of 4 lines of code I had to write just one simple line of code.

Example 2:

I will divide all the even numbers of the ‘digits’ list above and make a new list.

numbers = []for i in digits:
    if i%2 ==0:
        numbers.append(i/2)

print(numbers)

Output:

[0.0, 1.0, 2.0, 3.0, 4.0]

Here, we had to use an extra ‘if statement’ to check if the digit is even. Here is how to do the same thing in a list comprehension:

[i/2 for i in digits if i%2 ==0]

Output:

[0.0, 1.0, 2.0, 3.0, 4.0]

Example 3:

In this example, I will multiply the even elements of the ‘digits’ list by three and I will multiply the odd elements with three and add 1 to it to make it even. Here is the code:

numbers = []
for i in digits:
    if i%2 ==0:
        numbers.append(3*i)
    else:
        numbers.append(3*i + 1)
print(numbers)

Output:

[0, 4, 6, 10, 12, 16, 18, 22, 24, 28]

In this example, we have a for loop and an if-else statement. An implementation of the same code with list comprehension:

[3*i if i%2 ==0 else 3*i + 1 for i in digits]

Output:

[0, 4, 6, 10, 12, 16, 18, 22, 24, 28]

Example 4:

We can use the in-built functions in a list comprehension. I will import the math library and use the square root function in all the elements of the ‘digits’ list:

import math
[math.sqrt(i) for i in digits]

Output:

[0.0,
 1.0,
 1.4142135623730951,
 1.7320508075688772,
 2.0,
 2.23606797749979,
 2.449489742783178,
 2.6457513110645907,
 2.8284271247461903,
 3.0]

Example 5:

This example will show you how to use a custom function use it in each element of a list using list comprehension. First, let’s define a function and then use it on the ‘digits’ list:

def form(n):
    return 2**n - n**2[form(i) for i in digits]

Output:

[1, 1, 0, -1, 0, 7, 28, 79, 192, 431]

Look, how the formula in the function transforms each digit. I find it helpful to have the digit itself and the transformed number together. We can generate the tuples using the list comprehension where each tuple will include the digit itself and the transformation:

[(i, form(i)) for i in digits]

Output:

[(0, 1),
 (1, 1),
 (2, 0),
 (3, -1),
 (4, 0),
 (5, 7),
 (6, 28),
 (7, 79),
 (8, 192),
 (9, 431)]

Instead of tuples, we could use the sets as well:

[{i, form(i)} for i in digits]

Output:

[{0, 1},
 {1},
 {0, 2},
 {-1, 3},
 {0, 4},
 {5, 7},
 {6, 28},
 {7, 79},
 {8, 192},
 {9, 431}]

Example 6:

This example will work on a nested for loop. Let’s write a simple nested for loop and then we will see how to do that using list comprehension.

com = []for i in [3, 1, 4]:
    for j in [9, 0, 2]:
        com.append((i, j))
print(com)

Output:

[(3, 9), (3, 0), (3, 2), (1, 9), (1, 0), (1, 2), (4, 9), (4, 0), (4, 2)]

Here is the one-liner implementation using list comprehension of this nested for loop:

[(i, j) for i in [3, 1, 4] for j in [9, 0, 2]]

Output:

[(3, 9), (3, 0), (3, 2), (1, 9), (1, 0), (1, 2), (4, 9), (4, 0), (4, 2)]

Example 7:

I wanted to show an implementation of the if-else statement with nested for loop as well. This time I will check if the first list of Example 6 is an odd number. If it is odd, we will add 1 to it otherwise it will stay as it is.

com = []for i in [3, 1, 4]:
    for j in [9, 0, 2]:
        if i%2 != 0:
            com.append((i+1, j))
        else:
            com.append((i, j))

print(com)

Output:

[(4, 9), (4, 0), (4, 2), (2, 9), (2, 0), (2, 2), (4, 9), (4, 0), (4, 2)]

Below is the list comprehension implementation of the above code block:

[(i+1, j) if i%2 !=0 else (i, j) for i in [3, 1, 4] for j in [9, 0, 2]]

Output:

[(4, 9), (4, 0), (4, 2), (2, 9), (2, 0), (2, 2), (4, 9), (4, 0), (4, 2)]

Example 8:

In this example, I will use a 2-dimensional list and flatten it using a list comprehension approach.

v = [[1, 3, 4], [2, 5, 6], [7, 5, 4]][i for elements in v for i in elements]

Output:

[1, 3, 4, 2, 5, 6, 7, 5, 4]

Example 9:

We will take a 2-dimensional list for this example as well and rotate it to 90 degrees. This example is pretty popular in coding interviews. Let’s implement it with regular for loop first as it can be slightly complicated for new programmers:

td = [[1, 3, 4], 
      [2, 5, 6], 
      [7, 5, 4]]res = []for i in range(len(td)):
    el = []
    for j in td:
        el.append(j[i])
    res.append(el)

print(res)

Output:

[[1, 2, 7], [3, 5, 5], [4, 6, 4]]

Here is how to do it using list comprehension:

[[j[i] for j in td] for i in range(len(td))]

Output:

[[1, 2, 7], [3, 5, 5], [4, 6, 4]]

Conclusion

I hope these examples will be helpful for you to learn list comprehension well. While coding a big application, it is important to find ways to improve the code. List comprehension can be one of the ways to compact the codes and make them a bit faster as well. But list comprehension may not be a good idea for too complicated codes.

Feel free to follow me on Twitter and like my Facebook page.