List Comprehension

As it turns out, for loops in Python are quite slow. When you’re going through a list, there is a faster, more preferred way called list comprehension.

Let’s learn by example.

Filter

Imagine you want to filter out elements in a list. Let’s start with the old for-loop method.

original = [-10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

new_list = []

for item in original:
    if item > 0:
        new_list.append(item)

Wow, that was long. In Python, there’s almost always a really short way of doing things. Enter list comprehension.

new_list = [item for item in original if item > 0]

That’s it! One line of code! Woah. Let’s break it down.

One important point to remember is that list comprehensions always create new lists.

Let’s take another example.

Modification

Here, we want to go through a list and modify all the elements.

original = [-10, -9, -8, -7, -6, -5, -4, -3, -2, -1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

Again, let’s start with the for loop.

new_list = []
for item in original:
    if item % 2 == 0:
        new_list.append(item * item)
    else:
        new_list.append(0)

So what are we doing? If the number is even, we are squaring it. If it’s odd, we set it to be 0. How can we do this with list comprehension?

new_list = [x**2 if x % 2 == 0 else 0 for x in original]

Notice the slightly different syntax. Here, the if-else statement actually comes before the original. Why is this?

As it turns out, the if statement in Python has two purposes. One is the regular if statement that we all know and love; in this case, it goes after the original. Notably, if it goes after original, it cannot take an else clause.

The other purpose, however, is to serve as a substitute for the ternary operator. The ternary operator works like this:

a = 5 if 4 > 3 else 0

We can conditionally assign a a value depending on the if clause. In many languages, not Python though, the ternary operator has a different syntax. For example, in C and Java, it looks like this instead:

a = 4 > 3 ? 5 : 0

In Python, they decided not to use this syntax and instead just use the same old if-statement. Notably, the ternary operator must have an else clause in Python. And in this case, it goes _before the ordinary.

Printing

Printing is a rather special case. In Python 2, the version that we are using, print is not a function! Isn’t that weird? It’s a kind of strange statement. In future versions of Python, this has been fixed, and it is indeed a function once more.

Because it’s not a function, we cannot use print within list comprehension. The following will not work:

a = [1, 2, 3]
[print x for x in a]

Nope! It will throw a SyntaxError.

So how can we get around this? There are two ways.

from __future__ import print_function
a = [1, 2, 3]
[print(x) for x in a]

Two things have happened here. First, the import statement. This import statement makes print into a function!

Second, now that print is a function, we have to use function syntax.

# Before
print x

# After
print(x)

Note the differences.

The other way to get around this is simply to write your own function.

def my_print(stuff):
    print stuff

a = [1, 2, 3]
[my_print(x) for x in a]

And that, too, works.

Nested Loops

As it turns out, any set of nested loops that are iterating over lists can also be done with list comprehension!

x = '12345'
y = 'abc'

for num in x:
    for letter in y:
        print x+y

What does this print?

1a
1b
1c
2a
2b
2c
3a
3b
3c
4a
4b
4c
5a
5b
5c

Okay. How can we replicate this?

[print(x+y) for x in '12345' for y in 'abc']

Woah! Our list comprehensions can actually use multiple lists!

What about two dimensional lists?

original = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
for x in original:
    for y in x:
        print y

The output of this is simply

0
1
2
3
4
5
6
7
8

The list comprehension version?

[print(y) for x in original for y in x]

Pretty nifty!