List Comprehension

List comprehensions are used to derive a new list from an existing list. It's a short way to write one or multiple for loops.

A list comprehension consists of

  1. an input sequence

  2. an expression that yields the elements of the derived list

  3. an iterator over the input sequence

  4. a logical condition over the iterator (optional)

# Suppose you have a list of numbers and you want to get a corresponding list 
# with all the numbers multiplied by 2 only when the number itself is greater than 2. 
# List comprehensions are ideal for such situations.
listone = [2, 3, 4]
listtwo = [2*i for i in listone if i > 2]
print listtwo

List comprehension can support multiple levels of looping. But it's best practice to avoid more than two nested comprehensions and use conventional nested for loops.

# create product of two sets
A = [1, 2, 3]
B = ['a', 'b']
[(x, y) for x in A for y in B]
# [(1, 'a'), (1, 'b'), (2, 'a'), (2, 'b'), (3, 'a'), (3, 'b')]

# convert 2D list to 1D list
A = [[1, 2, 3], ['a', 'b', 'c']]
[x for y in A for x in y]   # first write outer for loop, then inside for loop
# [1, 2, 3, 'a', 'b', 'c']

# iterate over each entry in a 2D list
A = [[1, 2, 3], [4, 5, 6]]
[[x**2 for x in y] for y in A]
# [[1, 4, 9], [16, 25, 36]]

List comprehension also works on set.

Last updated