By Kevin Guan

2016-01-17 06:55:55 8 Comments

I have the following code:

[x**2 for x in range(10)]

When I run it in the Python Shell, it returns:

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

I've searched and it seems this is called a list comprehension, but how does it work?


@Matt Messersmith 2018-09-21 14:35:59

I've seen a lot of confusion lately (on other SO questions and from coworkers) about how list comprehensions work. A wee bit of math education can help with why the syntax is like this, and what list comprehensions really mean.

The syntax

It's best to think of list comprehensions as predicates over a set/collection, like we would in mathematics by using set builder notation. The notation actually feels pretty natural to me, because I hold an undergrad degree in Mathematics. But forget about me, Guido van Rossum (inventor of Python) holds a masters in Mathematics and has a math background.

Set builder notation crash course

Here's the (very basics) of how set builder notation works:

enter image description here

So, this set builder notation represents the set of numbers that are strictly positive (i.e. [1,2,3,4,...]).

Points of confusion

1) The predicate filter in set builder notation only specifies which items we want to keep, and list comprehension predicates do the same thing. You don't have to include special logic for omitting items, they are omitted unless included by the predicate. The empty predicate (i.e. no conditional at the end) includes all items in the given collection.

2) The predicate filter in set builder notation goes at the end, and similarly in list comprehensions. (some) Beginners think something like [x < 5 for x in range(10)] will give them the list [0,1,2,3,4], when in fact it outputs [True, True, True, True, True, False, False, False, False, False]. We get the output [True, True, True, True, True, False, False, False, False, False] because we asked Python to evaluate x < 5 for all items in range(10). No predicate implies that we get everything from the set (just like in set builder notation).

If you keep set builder notation in the back of your mind while using list comprehensions, they're a bit easier to swallow.


@Dave Rove 2018-09-11 09:00:15

If you prefer a more visual way of figuring out what's going on then maybe this will help:

# for the example in the question...

y = []
for x in range(10):
    y += [x**2]

# is equivalent to...

y = [x**2 for x in range(10)]

# for a slightly more complex example, it is useful
# to visualize  where the various x's end up...

a = [1,2,3,4]
b = [3,4,5,6]
c = []

for x in a:
          if x in b:
                  c += [x]
#   \         \        /
#    \    _____\______/
#     \  /      \
#      \/        \
#      /\         \
#     /  \         \
#    /    \         \
c = [x for x in a if x in b]


...produces the output [3, 4]

@cdlane 2016-01-21 19:21:39

There are list, dictionary, and set comprehensions, but no tuple comprehensions (though do explore "generator expressions").

They address the problem that traditional loops in Python are statements (don't return anything) not expressions which return a value.

They are not the solution to every problem and can be rewritten as traditional loops. They become awkward when state needs to be maintained & updated between iterations.

They typically consist of:

[<output expr> <loop expr <input expr>> <optional predicate expr>]

but can be twisted in lots of interesting and bizarre ways.

They can be analogous to the traditional map() and filter() operations which still exist in Python and continue to be used.

When done well, they have a high satisfaction quotient.

@bro-grammer 2018-01-09 10:37:08

This made my day: When done well, they have a high satisfaction quotient.

@Kevin Guan 2016-01-17 06:55:55

From the documentation:

List comprehensions provide a concise way to create lists. Common applications are to make new lists where each element is the result of some operations applied to each member of another sequence or iterable, or to create a subsequence of those elements that satisfy a certain condition.

About your question, the list comprehension does the same thing as the following "plain" Python code:

>>> l = [] 
>>> for x in range(10):
...     l.append(x**2)
>>> l
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

How do you write it in one line? Hmm...we can...probably...use map() with lambda:

>>> list(map(lambda x: x**2, range(10)))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

But isn't it clearer and simpler to just use a list comprehension?

>>> [x**2 for x in range(10)]
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

Basically, we can do anything with x. Not only x**2. For example, run a method of x:

>>> [x.strip() for x in ('foo\n', 'bar\n', 'baz\n')]
['foo', 'bar', 'baz']

Or use x as another function's argument:

>>> [int(x) for x in ('1', '2', '3')]
[1, 2, 3]

We can also, for example, use x as the key of a dict object. Let's see:

>>> d = {'foo': '10', 'bar': '20', 'baz': '30'}
>>> [d[x] for x in ['foo', 'baz']]
['10', '30']

How about a combination?

>>> d = {'foo': '10', 'bar': '20', 'baz': '30'}
>>> [int(d[x].rstrip('0')) for x in ['foo', 'baz']]
[1, 3]

And so on.

You can also use if or if...else in a list comprehension. For example, you only want odd numbers in range(10). You can do:

>>> l = []
>>> for x in range(10):
...     if x%2:
...         l.append(x)
>>> l
[1, 3, 5, 7, 9]

Ah that's too complex. What about the following version?

>>> [x for x in range(10) if x%2]
[1, 3, 5, 7, 9]

To use an if...else ternary expression, you need put the if ... else ... after x, not after range(10):

>>> [i if i%2 != 0 else None for i in range(10)]
[None, 1, None, 3, None, 5, None, 7, None, 9]

Have you heard about nested list comprehension? You can put two or more fors in one list comprehension. For example:

>>> [i for x in [[1, 2, 3], [4, 5, 6]] for i in x]
[1, 2, 3, 4, 5, 6]

>>> [j for x in [[[1, 2], [3]], [[4, 5], [6]]] for i in x for j in i]
[1, 2, 3, 4, 5, 6]

Let's talk about the first part, for x in [[1, 2, 3], [4, 5, 6]] which gives [1, 2, 3] and [4, 5, 6]. Then, for i in x gives 1, 2, 3 and 4, 5, 6.

Warning: You always need put for x in [[1, 2, 3], [4, 5, 6]] before for i in x:

>>> [j for j in x for x in [[1, 2, 3], [4, 5, 6]]]
Traceback (most recent call last):
  File "<input>", line 1, in <module>
NameError: name 'x' is not defined

We also have set comprehensions, dict comprehensions, and generator expressions.

set comprehensions and list comprehensions are basically the same, but the former returns a set instead of a list:

>>> {x for x in [1, 1, 2, 3, 3, 1]}
{1, 2, 3}

It's the same as:

>>> set([i for i in [1, 1, 2, 3, 3, 1]])
{1, 2, 3}

A dict comprehension looks like a set comprehension, but it uses {key: value for key, value in ...} or {i: i for i in ...} instead of {i for i in ...}.

For example:

>>> {i: i**2 for i in range(5)}
{0: 0, 1: 1, 2: 4, 3: 9, 4: 16}

And it equals:

>>> d = {}
>>> for i in range(5):
...     d[i] = i**2
>>> d
{0: 0, 1: 1, 2: 4, 3: 9, 4: 16}

Does (i for i in range(5)) give a tuple? No!, it's a generator expression. Which returns a generator:

>>> (i for i in range(5))
<generator object <genexpr> at 0x7f52703fbca8>

It's the same as:

>>> def gen():
...     for i in range(5):
...         yield i
>>> gen()
<generator object gen at 0x7f5270380db0>

And you can use it as a generator:

>>> gen = (i for i in range(5))
>>> next(gen)
>>> next(gen)
>>> list(gen)
[2, 3, 4]
>>> next(gen)
Traceback (most recent call last):
  File "<input>", line 1, in <module>

Note: If you use a list comprehension inside a function, you don't need the [] if that function could loop over a generator. For example, sum():

>>> sum(i**2 for i in range(5))

Related (about generators): Understanding Generators in Python.

@AChampion 2016-01-17 07:11:45

The ternary expression x if cond else y really doesn't have anything specifically to do with list comprehensions - any valid expression can be used inside a list compression - the ternary expression is one of many python expressions.

@Kevin Guan 2016-01-17 07:14:40

@AChampion: Yeah, I mentioned that in my answer because I tried [i for i in x if i else y] when I was learning list comprehensions and it doesn't work. After some research I understand that I must use [i if i else y for i in x] instead. So I think if I mention it here then others can avoid the issue which I had before.

@styvane 2016-01-21 19:54:46

It is worth to mention that you only need to use the list constructor: list(<map object>) in Python 3.x

@Kevin Guan 2016-01-22 01:05:26

@user3100115: Yep, I'm using it in my examples. Check that list(map(lambda x: x**2, range(10))) one.

@midori 2016-02-02 02:19:17

it's a bad practise to put 2 or more for in on list comprehension

@Kevin Guan 2016-02-02 04:18:46

@minitoto: What's wrong? It's a feature of list comprehension called nested list comprehension.

@midori 2016-02-02 04:20:10

it becomes less readable

@Kevin Guan 2016-02-02 04:37:35

@minitoto: Well, this answer is more like a "List comprehension FAQ". So I think I should mention all the features of list comprehension here. And I think it's not so unreadable to me, but however, if you don't like it, sure you can also use a for loop instead of that.

Related Questions

Sponsored Content

10 Answered Questions

[SOLVED] Iterating over dictionaries using 'for' loops

19 Answered Questions

[SOLVED] How do I remove an element from a list by index in Python?

  • 2009-03-09 18:16:11
  • Joan Venge
  • 2193091 View
  • 1230 Score
  • 19 Answer
  • Tags:   python list

30 Answered Questions

[SOLVED] How do I check if a list is empty?

  • 2008-09-10 06:20:11
  • Ray Vega
  • 2217118 View
  • 3237 Score
  • 30 Answer
  • Tags:   python list

22 Answered Questions

[SOLVED] How do I list all files of a directory?

  • 2010-07-08 19:31:22
  • duhhunjonn
  • 3217816 View
  • 3475 Score
  • 22 Answer
  • Tags:   python directory

34 Answered Questions

[SOLVED] How do I sort a dictionary by value?

28 Answered Questions

[SOLVED] How do I concatenate two lists in Python?

31 Answered Questions

[SOLVED] "Least Astonishment" and the Mutable Default Argument

38 Answered Questions

[SOLVED] What does the "yield" keyword do?

17 Answered Questions

[SOLVED] How to make a chain of function decorators?

27 Answered Questions

[SOLVED] What does if __name__ == "__main__": do?

Sponsored Content