Nested List Comprehensions in Python

07 Nov, 2018
List Comprehensions are one of the most amazing features of Python. It is a smart and concise way of creating lists by iterating over an iterable object. Nested List Comprehensions are nothing but a list comprehension within another list comprehension which is quite similar to nested for loops.
Let’s take a look at some examples to understand what nested list comprehensions can do:
Example 1:
I want to create a matrix which looks like below:

matrix = [[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]]

The below code uses nested for loops for the given task:
matrix = []  for i in range(5):          # Append an empty sublist inside the list    matrix.append([])          for j in range(5):        matrix[i].append(j)          print(matrix)
[[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]]

The same output can be achieved using nested list comprehension in just one line:
# Nested list comprehensionmatrix = [[j for j in range(5)] for i in range(5)]  print(matrix)
[[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]]

Explanation:
The syntax of the above program is shown below:
[expression for i in range(5)] –> which means that execute this expression and append its output to the list until variable i iterates from 0 to 4.
For example:- [i for i in range(5)] –> In this case, the output of the expressionis simply the variable i itself and hence we append its output to the list while iiterates from 0 to 4.
Thus the output would be –> [0, 1, 2, 3, 4]
But in our case, the expression itself is a list comprehension. Hence we need to firstsolve the expression and then append its output to the list.
expression = [j for j in range(5)] –> The output of this expression is same as theexample discussed above.
Hence expression = [0, 1, 2, 3, 4].
Now we just simply append this output until variable i iterates from 0 to 4 which wouldbe total 5 iterations. Hence the final output would just be a list of the output of theabove expression repeated 5 times.
Output: [[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]]
Example 2:
Suppose I want to flatten a given 2-D list:

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

Expected Output: flatten_matrix = [1, 2, 3, 4, 5, 6, 7, 8, 9]

This can be done using nested for loops as follows:
# 2-D Listmatrix = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]  flatten_matrix = []  for sublist in matrix:    for val in sublist:        flatten_matrix.append(val)          print(flatten_matrix)
[1, 2, 3, 4, 5, 6, 7, 8, 9]

Again this can be done using nested list comprehension which has been shown below:
# 2-D Listmatrix = [[1, 2, 3], [4, 5], [6, 7, 8, 9]]  # Nested List Comprehension to flatten a given 2-D matrixflatten_matrix = [val for sublist in matrix for val in sublist]  print(flatten_matrix)
[1, 2, 3, 4, 5, 6, 7, 8, 9]

Explanation:
In this case, we need to loop over each element in the given 2-D list and append itto another list. For better understanding, we can divide the list comprehension intothree parts:
flatten_matrix = [val
                  for sublist in matrix
                  for val in sublist]
The first line suggests what we want to append to the list. The second line is theouter loop and the third line is the inner loop.
‘for sublist in matrix’ returns the sublists inside the matrix one by one which would be:
[1, 2, 3], [4, 5], [6, 7, 8, 9]
‘for val in sublist’ returns all the values inside the sublist.
Hence if sublist = [1, 2, 3], ‘for val in sublist’ –> gives 1, 2, 3 as output one by one.
For every such val, we get the output as val and we append it to the list.
Example 3:
Suppose I want to flatten a given 2-D list and only include those strings whose lengths are less than 6:
planets = [[‘Mercury’, ‘Venus’, ‘Earth’], [‘Mars’, ‘Jupiter’, ‘Saturn’], [‘Uranus’, ‘Neptune’, ‘Pluto’]]
Expected Output: flatten_planets = [‘Venus’, ‘Earth’, ‘Mars’, ‘Pluto’]
This can be done using an if condition inside a nested for loop which is shown below:
# 2-D List of planetsplanets = [['Mercury', 'Venus', 'Earth'], ['Mars', 'Jupiter', 'Saturn'], ['Uranus', 'Neptune', 'Pluto']]  flatten_planets = []  for sublist in planets:    for planet in sublist:                  if len(planet) < 6:            flatten_planets.append(planet)          print(flatten_planets)
['Venus', 'Earth', 'Mars', 'Pluto']

This can also be done using nested list comprehensions which has been shown below:
# 2-D List of planetsplanets = [['Mercury', 'Venus', 'Earth'], ['Mars', 'Jupiter', 'Saturn'], ['Uranus', 'Neptune', 'Pluto']]  # Nested List comprehension with an if conditionflatten_planets = [planet for sublist in planets for planet in sublist if len(planet) < 6]          print(flatten_planets)
['Venus', 'Earth', 'Mars', 'Pluto']

Explanation:
This example is quite similar to the previous example but in this example, we justneed an extra if condition to check if the length of a particular planet is less than6 or not.
This can be divided into 4 parts as follows:
flatten_planets = [planet 
                   for sublist in planets 
                   for planet in sublist 
                   if len(planet) < 6] 

Picked
python-list
Technical Scripter 2018
Python
Technical Scripter
python-list
Writing code in comment?
Please use ide.geeksforgeeks.org,
generate link and share the link here.