NB: Nested Comprehensions

Programming for Data Science

Recall that nested loops are the algorithmic complement to nested data structures.

Just as we can nest loops using for loops, we can do so with comprehensions.

Here are some examples.

This notebook is adapted from GeeksForGeeks.

Example 1: Creating a Matrix

Here is how we can make a matrix — a two-dimensional data structure where each element is of the same data type — using plain old for loops.

matrix1 = [] # Matrix created
for i in range(5):
    matrix1.append([]) # Row created
    for j in range(5):
        matrix1[i].append(j) # Cell populated

matrix1

[[0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4]]

Here’s how we can do this with a nested list comprehension:

matrix2 = [[j for j in range(5)] for i in range(5)]

matrix2

[[0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4],
 [0, 1, 2, 3, 4]]

We replace five lines of code with one.

Example 2: Filtering a Nested List

Let create a filter to extract some things we want from a matrix.

In this case, we want to pull out all of the odd numbers and save them as a list.

matrix3 = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
odd_numbers1 = []
for row in matrix3:
    for element in row:
        if element % 2 != 0:
            odd_numbers1.append(element)

odd_numbers1

[1, 3, 5, 7, 9]

odd_numbers2 = [element for row in matrix3 for element in row if element % 2 != 0]

odd_numbers2

[1, 3, 5, 7, 9]

Here’s a closer look at how nesting is represented in a comprehension compared to the traditional way.

Traditional for loop:

for row in matrix3:
    for element in row:
        if element % 2 != 0:
            odd_numbers1.append(element) # This is hoised to the top
                                         # and implicitly appended

Comprehension (inside []):

element                        
    for row in matrix3 
        for element in row 
            if element % 2 != 0

See how they are the same, except that the kernel operation — appending filtered values to a result list — is hoisted to the top in the case of the comprehension.

Example 3: Flattening Nested Sub-Lists

Here we take a nested structure and flatten it out.

nested_list = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]
flat_list1 = []
for sub_list in nested_list:
    for val in sub_list:
        flat_list1.append(val)

flat_list1

[1, 2, 3, 4, 5, 6, 7, 8, 9]

flat_list2 = [val for sub_list in nested_list for val in sub_list]

flat_list2

[1, 2, 3, 4, 5, 6, 7, 8, 9]

Example 4: Manipulating Matrix Values

Finally, we demonstrate how to manipulate each value in a matrix of words.

In this case, we simply capitalize each string.

matrix4 = [["apple", "banana", "cherry"],
        ["date", "fig", "grape"],
        ["kiwi", "lemon", "mango"]]
modified_matrix1 = []
for row in matrix4:
    modified_row = []
    for fruit in row:
        modified_row.append(fruit.capitalize())
    modified_matrix1.append(modified_row)

modified_matrix1

[['Apple', 'Banana', 'Cherry'],
 ['Date', 'Fig', 'Grape'],
 ['Kiwi', 'Lemon', 'Mango']]

modified_matrix2 = [[fruit.capitalize() for fruit in row] for row in matrix4]

modified_matrix2

[['Apple', 'Banana', 'Cherry'],
 ['Date', 'Fig', 'Grape'],
 ['Kiwi', 'Lemon', 'Mango']]