List Comprehensions in Python

Lists are a versatile Python data structure and allows various sequence operations. Besides the usual, Lists also include a powerful operation called as ‘list comprehension’ which allows us to process various list structures, like a matrix or nested arrays or lists. It allows you to concisely form a new list by filtering the elements of a list.

List comprehension takes the following form.

[expr for val in collection if condition]

Which is equivalent to the following for loop.

result = []
for val in collection:
   if condition:

Take the example matrix given below.

>>> matrix = [[1,2,3],
>>> matrix
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]

It is easy to get rows from the matrix as they stored in a row format, like given below.

>>> matrix[0]
[1, 2, 3]
>>> matrix[1]
[4, 5, 6]
>>> matrix[2]
[7, 8, 9]

But what if you wanted to grab the columns from the matrix. Here we can use list comprehensions. List comprehensions are a way to build new lists by running a expression on each element in the sequence.

>>> col1 = [row[0] for row in matrix]
>>> col1
[1, 4, 7]
>>> col2 = [row[1] for row in matrix]
>>> col2
[2, 5, 8]
>>> col3 = [row[2] for row in matrix]
>>> col3
[3, 6, 9]

The first line above basically asks to give us the item at location row[0] (first item in the row) for each row in the matrix which is than stored in the new ‘col1’ list. So for example on the first iteration of the matrix row[0] will return ‘1’, on the second iteration it will return ‘4’ and lastly ‘7’.

We can also make list comprehensions more complex.

# Here we have added 1 to each column element before storing in the new list
>>> col1_new = [row[0] + 1 for row in matrix]  
>>> col1_new
[2, 5, 8]
# Double each character found in the string and create a new list.
>>> [c * 2 for c in 'Hello']
['HH', 'ee', 'll', 'll', 'oo']

Note that list comprehensions has to be in square brackets which basically hints that a list is returned. The following for example will not work.

>>> col1 = row[0] for row in matrix
  File "<stdin>", line 1
    col1 = row[0] for row in matrix
SyntaxError: invalid syntax

You can also add a filter condition to the expression.

# Get only those elements which are greater than 1
>>> [row[0] for row in matrix if row[0] > 1]
[4, 7]
# Get only numbers which are even
>>> [row[1] for row in matrix if row[1] % 2 == 0]
[2, 8]
>>> [row[0] for row in matrix if row[0] % 2 == 0]

This is just a introductory article around list comprehensions, as they are more powerful in what they can accomplish. We will cover them more in future articles.

Leave a Reply

Your email address will not be published. Required fields are marked *