Using List Comprehension in Python

Python List Comprehension

Hello everyone! Today, we’ll talk about how we can use List Comprehension in Python.

List Comprehension is generally a syntactic sugar to make code easier to read and write.

Often, when we deal with code involving creating lists, it is cumbersome to write nested loops again and again.

Python has made this easier for us by introducing this feature.

Let’s now take a look to understand how we can use this in our programs, by taking suitable examples!


Basic Structure of a List Comprehension

Let’s consider the following code, written normally:

word = "Hello from AskPython"
letters = []

for letter in word:
    letters.append(letter)

print(letters)

Output

['H', 'e', 'l', 'l', 'o', ' ', 'f', 'r', 'o', 'm', ' ', 'A', 's', 'k', 'P', 'y', 't', 'h', 'o', 'n']

The above snippet prints a list of the letters in our word.

We can use List Comprehension to shorten this code, since the elements of the list have a common property: they are letters, and they will be appended to the list.

Let’s now use List Comprehension to make this more shorter and readable:

word = "Hello from AskPython"

letters = [letter for letter in word]

print(letters)

Output

['H', 'e', 'l', 'l', 'o', ' ', 'f', 'r', 'o', 'm', ' ', 'A', 's', 'k', 'P', 'y', 't', 'h', 'o', 'n']

See how easy it is? The intent of the code is made clear: we pick up letters of the word and add it to our list directly!

Now, we can use list comprehension with other iterables too!

Let’s take another example, where we can generate the squares of numbers from 1 to 10.

The normal approach will be the following:

squares = []

for i in range(1, 11):
    squares.append(i * i)

print(squares)

Output

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

While this is small enough, we can do better using list comprehension. We can reduce this to just two lines of code!

squares = [i * i for i in range(1, 11)]
print(squares)

We’ve now showed you the power of List Comprehension! Let’s now add some more spice to it, by dealing with conditionals like if and else!


Using Conditions in a List Comprehension

We can use if and else conditionals in our List Comprehension.

Let’s consider the first case, where we only have an if condition.

The general structure of this type of list comprehension is as follows:

list = [item for item in iterable if condition]

So here, list will only consist of items where condition holds True.

Let’s take our previous example of constructing squares, and restrict it to only even numbers, using if.

squares = [i * i for i in range(1, 11) if i % 2 == 0]
print(squares)

Output

[4, 16, 36, 64, 100]

Indeed, we can only get the even element squares here, since i % 2 == 0 only if i is even.

Let’s now take the second case, where we have an else condition as well. The structure will now look like this:

list = [value1 if condition else value2 for item in iterable]

Here, the list will comprise elements of value1 if condition == True and elements of value2 if condition == False.

Let’s now take an example where we keep printing integer squares until i<=5. If i > 5, we’ll print 0 instead.

Our list comprehension will now look like this:

my_list = [i * i if i <= 5 else 0 for i in range(10)]
print(my_list)

Output

[0, 1, 4, 9, 16, 25, 0, 0, 0, 0]

As you can see, the list contains the squares of all numbers <= 5 only. The remaining elements are set to 0!

We can also use other conditionals, and even lambda functions, if we want to!

Here is a slightly contrived example, which uses a lambda to compute the successive pair-sums from 0. (0, 1 + 2, 2 + 3, 3 + 4.. )

pair_sums = [(lambda x, y: x + y)(i, j) if i > 1 else 0 for i, j in zip(range(1, 11), range(0, 10))]
print(pair_sums)

Output

[0, 3, 5, 7, 9, 11, 13, 15, 17, 19]

As you can see, this code is not the most readable, and you might have been better off with something else!

So, be careful not to use List Comprehensions if you’re trying to do too much at once. It’s better to stick with this method when you want to run a simple loop-conditional statement to build lists, not when you want to perform mathematical calculations on each element individually.


Conclusion

In this article, we learned about using Python’s List Comprehension semantics. This makes it easier to reduce writing repetitive code again and again using iterative loops!

References