Python: List Comprehensions

Note: Lines beginning with ">>>" and "..." indicate input to Python (these are the default prompts of the interactive interpreter). Everything else is output from Python.

Python supports a concept called "list comprehensions". It can be used to construct lists in a very natural, easy way, like a mathematician is used to do.

The following are common ways to describe lists (or sets, or tuples, or vectors) in mathematics.

S = {x : x in {0 ... 9}}
V = (1, 2, 4, 8, ..., 2)
M = {x | x in S and x even}

You probably know things like the above from mathematics lessons at school. In Python, you can write these expression almost exactly like a mathematician would do, without having to remember any special cryptic syntax.

This is how you do the above in Python:

>>> S = [x**2 for x in range(10)]
>>> V = [2**i for i in range(13)]
>>> M = [x for x in S if x % 2 == 0]
>>> 
>>> print S; print V; print M
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
[0, 4, 16, 36, 64]

I'm sure you want to see a more complicated example. :-) The following is yet another way to compute prime numbers. The interesting thing is that we first build a list of non-prime numbers, using a single list comprehension, then use another list comprehension to get the "inverse" of the list, which are prime numbers.

>>> noprimes = [j for i in range(2, 8) for j in range(i*2, 50, i)]
>>> primes = [x for x in range(2, 50) if x not in noprimes]
>>> print primes
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47]

NB: You can nest list comprehensions inside of each other, so you could write the above example with a single statement (without the need for the temporary variable "noprimes"). However, such lines tend to get long and less readable, so this is not recommended.

Of course, list comprehensions don't only work for numbers. Lists can contain any type of elements, including strings, nested lists and functions. You can even mix different types within a list.

The following works on a list of strings and produces a list of lists. Each of the sublists contains two strings and an integer.

>>> words = 'The quick brown fox jumps over the lazy dog'.split()
>>> print words
['The', 'quick', 'brown', 'fox', 'jumps', 'over', 'the', 'lazy', 'dog']
>>> 
>>> stuff = [[w.upper(), w.lower(), len(w)] for w in words]
>>> for i in stuff:
...     print i
... 
['THE', 'the', 3]
['QUICK', 'quick', 5]
['BROWN', 'brown', 5]
['FOX', 'fox', 3]
['JUMPS', 'jumps', 5]
['OVER', 'over', 4]
['THE', 'the', 3]
['LAZY', 'lazy', 4]
['DOG', 'dog', 3]
>>> 
>>> stuff = map(lambda w: [w.upper(), w.lower(), len(w)], words)
>>> for i in stuff:
...     print i
... 
['THE', 'the', 3]
['QUICK', 'quick', 5]
['BROWN', 'brown', 5]
['FOX', 'fox', 3]
['JUMPS', 'jumps', 5]
['OVER', 'over', 4]
['THE', 'the', 3]
['LAZY', 'lazy', 4]
['DOG', 'dog', 3]

The above example also demonstrates that you can do exactly the same thing with map() and a lambda function. However, there are cases when you cannot use map() and have to use a list comprehension instead, or vice versa. When you can use both, then it is often preferable to use a list comprehension, because this is more efficient and easier to read, most of the time.

You cannot use list comprehensions when the construction rule is too complicated to be expressed with "for" and "if" statements, or if the construction rule can change dynamically at runtime. In this case, you better use map() and / or filter() with an appropriate function. Of course, you can combine that with list comprehensions.


[HTML 4.01]