12.0 List Comprehensions (optional)¶
Python has a useful and powerful syntax for generating lists. A list comprehension in Python is a construct for creating a list based on another iterable structure (e.g., a list, a sequence, etc) in a single line of code.
For example, given a list of numbers, 'xlist' (say), a list of squares of those numbers may be generated as follows:
xlist = [1, 2, 3, 4, 5, 6]
x2list = [x**2 for x in xlist] # basic list comprehension
print(x2list) # check out 'x2list'....
# the above lines of code are equivalent to:
xlist = [1, 2, 3, 4, 5, 6]
x2list = [] # initialize empty list
for x in xlist:
x2list.append(x**2)
print(x2list)
List comprehensions can also contain conditional statements:
# let's say you want to square out only the odd values:
xlist = list(range(1,7)) # same list as before
x2list_new = [x**2 for x in xlist if x % 2] # consider only ODD values of x
print(x2list_new)
In the above conditional we have used $x\%2$, which is the remainder of the division of $x$ by $2$; the result is either $0$ (i.e., 'False') or $1$ (i.e., 'True'). Hence, if the number $x$ is divisible by $2$ the conditional evaluates to 'False' and we do not square out that particular $x$. On the other hand, for odd values of $x$ the conditional evaluates to 'True', and so the value will be squared out.
If you require more complex transformations on the elements of the original list, the 'if-else' expression must appear before the 'for-loop':
# check this out:
xlist = list(range(1,7))
x23list = [x**2 if x%2 else x**3 for x in xlist]
print(xlist)
print(x23list)
# what happened?
At this point we can state the general design pattern for these list comprehensions:
<your_list_name> = [<expression(x)> for x in <list> if <conditional>]
# which is equivalent to:
for x in <list>: # <list> is your own sequence
if <conditional>: # <conditional> involves some conditions you may want to test
<expression(x)> # an expression that depends on 'x'
You can also use list comprehensions with strings:
text = "Keep calm and eat ice cream".split()
# a first example:
list1 = [word.upper() for word in text]
print(list1)
# a second example:
list2 = [len(word) for word in text]
print(list2)
List comprehensions can be nested. The general design pattern is:
[<expression> for <outer_var> in <outer_seq> for <inner_var> in <inner_seq>]
This looks a bit complicated.... Here is an example:
coords = [(x,y) for x in range(5) for y in range(3)]
print(coords)
# let's prints these points:
import numpy as np
from matplotlib import pyplot as plt
for item in coords:
plt.plot(item[0], item[1],'bo')
plt.xlabel('x'); plt.ylabel('y')
plt.grid(); plt.show()
# you could write the 'coords' list comprehension as follows:
mylist = [] # initialize storage
for x in range(5): # outer loop (range(5) is the <outer_seq>)
for y in range(3): # inner loop (range(3) is the <inner_seq>)
mylist.append((x,y)) # appending one tuple at a time
print(mylist) # same thing as above....
You can use the zip() function in list comprehensions:
# the list below must have the same length
a = [1, 2, 3] # an arbitrary list
b = [4, 5, 6] # another list
# example 1:
list1 = [(x,y) for x, y in zip(a,b)]
print('list1 = ', list1)
# remember that you can't print directly the output of 'zip()'
# example 2:
list2 = [x*y for x, y in zip(a,b)]
print('list2 = ', list2)
You can "flatten" two-dimensional arrays using nested list comprehensions:
a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
[item for row in a for item in row]
# a more normal way to write the above:
mylist = [] # initialize storage (1D list)
for row in a: # outer loop
for item in row: # inner loop
mylist.append(item) # update the one-dimensional list
print(mylist) # display one-dimensional list
# KEY POINT: the outer loop comes before the inner loop