List comprehensions

Mathematicians will know that dealing with vectors, matrices, operators, etc. involves a lot of repetitive and iterative calculations. Programmers will all know that a for loop is the natural way to deal with such a procedure. For example, let’s say we want a list with some particular numbers in it:

squares = []
for i in range(8):

This will create a list of all square numbers up to but not including \(8^2\). This is fine, but in my opinion Python has a much more natural way to work here, and that is the notion of a list comprehension. Using a list comprehension, we write

squares = [i**2 for i in range(8)]

This gives the same result in a single line. Not a huge gain, I’ll grant you, but it would still be writable as a single line whether we had three or thirty for loops (definitely not recommended though). As well as that, it would still be a single line if we had three or thirty for loops and three or thirty if/elif/else checks to boot. But even disregarding all that, I think it has the advantage because it is much closer to a construct that any mathematician will recognise. Our list comprehension is identical to $$\mathrm{squares} = \{a^2 : 0\leq a \lt 8\}.$$ This set notation is ubiquitous in mathematics and if you align the colon with the word for you will immediately see the perfect mapping between it and the list comprehension. It will almost always be possible to replace for loops with comprehensions.

To drive home the point I will write a function to return the sum of two matrices, once in loop form and once in comprehension form (using the rows() and columns() functions from the previous page).

def add(matrix1, matrix2):
    row_range = range(rows(matrix))
    col_range = range(columns(matrix))
    matrix = list()
    for i in row_range:
        for j in col_range:
            matrix[i].append(matrix1[i][j] + matrix2[i][j])
    return matrix

Not bad. I don’t think there are any major savings to be made if you want to stick to using a for loop. As a list comprehension:

def add(matrix1, matrix2):
    row_range = range(rows(matrix1))
    col_range = range(columns(matrix1))
    return [[matrix1[i][j] + matrix2[i][j] for j in col_range] for i in row_range]

Much nicer. Potentially a little harder to parse at first, but much briefer and after a few it will seem very natural. The inner loop here corresponds to the second loop in the previous example, and the outer loop to the top level loop. This holds true in general for any number of loops. Top-down for the traditional loops reads as out-in for the list comprehensions.


So far our functional repertoire is:

  • rows(matrix): count the number of rows in a matrix.
  • columns(matrix): count the number of columns in a matrix.
  • add(matrix1, matrix2): compute the element-wise sum of two matrices.

Previous: Preliminaries
Next: Some stock functions

Leave a Reply