Section 2: Nested List Comprehensions

 Table of Contents > Chapter 4 > Section 2 

We are often in situations where we deal with information in the world that spans two or more dimensions.  For example, atmospheric models predict air pressure and wind velocities over three-dimensional regions.  Engineering models predict the deflection of a two-dimensional plate when subjected to a certain load.

Information that is organized in a two-dimensional grid can be represented as a list of lists.  There is more than one way that this representation can occur.  We can split the two-dimensional grid into rows and represent information on each row with a list, or we can split the two-dimensional grid into columns.  The choice is somewhat arbitrary and may depend on the particular problem that we are trying to solve.

Example

Suppose we have a two-dimensional plate aligned with the x and y axes as illustrated in Figure 4.1.  Further suppose that physical models predict that when the plate is heated in a certain way, the temperature at a point (x, y) on the plate is given by the formula:

T(x, y) = 20 * (x² + y²)

Now suppose we want to know the temperature at all points on a two-dimensional grid as illustrated in Figure 4.1 below. 

This information can be represented as a list of lists.  We can think of the problem this way:

lox = [2, 4, 6, 10, 12]

gridTemps = []

for x in lox:

     <compute column of temperatures - colTemps>

     gridTemps = gridTemps + [colTemps]
Now, to create a column of temperatures for a given value of x:

colTemps = []

for y in [2, 4, 7, 9]:

     colTemps = colTemps + [T(x, y)]
Note that we've assumed the existence of a Python function T that, given the x and y coordinates of a point on the grid, produces the corresponding temperature at that point.

Putting it all together we have:

lox = [2, 4, 6, 10, 12]

loy = [2, 4, 7, 9]

gridTemps = []

for x in lox:

     colTemps = []

     for y in loy:

          colTemps = colTemps + [T(x, y)]

     gridTemps = gridTemps + [colTemps]
Now we can start re-factoring this code using list comprehensions.  First, we work on the inner loop:

lox = [2, 4, 6, 10, 12]

loy = [2, 4, 7, 9]

gridTemps = []

for x in lox:

     colTemps = [T(x, y) for y in loy]

     gridTemps = gridTemps + [colTemps]
then the outer loop:

lox = [2, 4, 6, 10, 12]

loy = [2, 4, 7, 9]

gridTemps = 
[[T(x, y) for y in loy] for x in lox]

Assuming that T(x, y) is computed using the function given earlier, gridTemps is assigned the value:

[[8, 20, 53, 85], 

 [20, 32, 65, 97], 

 [40, 52, 85, 117], 

 [104, 116, 149, 181], 

 [148, 160, 193, 225]]

Note that the type of gridTemps is (listof (listof int)).  Further note that gridTemps[0] is the list [8, 20, 53, 85] and that this represents temperatures at the points (2, 2), (2, 4), 
(2, 7) and (2, 9) - in other words, at those points lying in the first column of the grid.