How to Replace a For Loop With List Comprehension in a 2D Matrix?

Replace For Loop With List Comprehension (1)

A For Loop is an iterator used to iterate over a sequence of objects like strings, lists, tuples, and sets. It allows us to execute a set of statements over and over again for the items of the sequence.

We can use the for keyword to implement a for loop.

Refer to this post to know more about for loop.

Apart from iterating over certain objects, for loops are used for other purposes like creating a nested dictionary, appending multiple pandas data frames into one, etc.

Another similar use case is How to create a countdown with for loops.

For loop also plays a significant role in creating matrices, accessing the elements of the matrix, adding the elements, and so on. While using for loop might seem an easy way to work with matrices, it might become lengthy if you are dealing with high dimensional matrices.

For simplicity and readability, there are other methods to carry out the same task with more efficiency such as list comprehensions.

Learn about list comprehensions here.

Replace a For Loop With List Comprehension in a 2D Matrix

In this section, we are going to look at all the possible use cases of for loops in a 2D matrix and also the alternative approach to it.

A matrix is an essential data storage in the field of statistics, data processing, image processing, etc

A 2D matrix is also known as a 2D array and consists of two dimensions – a row and a column to store the elements. We can implement a Python Matrix in the form of a list or an array.

We can also conveniently use the numpy library to create a matrix.

Refer to this post to learn how to create a matrix using numpy.

Accessing the Elements of a Matrix

Once the matrix is created, you’d want to print the elements separately. We can use for loop to iterate over every row and print the row elements.

import numpy as np
mat= np.array([[ 10, 20, 30],[ 40, 50, 60]])
print("The 2D matrix is:")
print(mat)
print("The elements of the matrix are:")
for row in mat:
    for element in row:
        print(element)

We are creating a 2D matrix using the numpy library’s array method. The matrix is made up of 2 rows and 3 columns and contains 6 elements. This matrix is named mat. It is printed in the next line.

We are using a for loop to get through every row in the matrix and every element in the row. The element at each position is printed.

Accessing Elements With For Loop
Accessing Elements With For Loop

The list comprehension follows the same code but doesn’t need the code spread across multiple lines. We can do the same thing in just a single line of code. The elements of the matrix when accessed by list comprehension are stored in the form of a list.

element = [ele for row in mat for ele in row]
print("The elements of the matrix are:\n",element)

Instead of going with a loop, we are creating an iterator called ele that accesses the elements of every row within the matrix.

Accessing Elements With List Comprehension
Accessing Elements With List Comprehension

Modifying the Elements of a Matrix

We can even modify the elements of a matrix with the help of for loop but it is a tad bit lengthy process.

matA = [[1, 2, 3],
          [4, 5, 6]]
print("The original matrix is :\n",matA)
for i in range(len(matA)):
    for j in range(len(matA[i])):
        if matA[i][j] % 2 == 0:
            matA[i][j] *= 3
print("The modified matrix is:\n",matA)

We are creating a matrix called matA with 2 rows and 3 columns. This matrix is printed in the next line.

We are initializing two for loops. The outer loop is to run across the row elements and the inner loop is to run through the columns. The element’s position is given by [i][j]. If the current element is divisible by two, the element is multiplied by three. The other elements remain untouched.

Modifying The Elements Using For Loop
Modifying The Elements Using For Loop

Let us see how to achieve the same results using list comprehension.

matA= [[1, 2, 3],
          [4, 5, 6]]
print("The original matrix is :\n",matA)
# Modifying elements using list comprehensions
matM= [[element * 3 if element % 2 == 0 else element for element in row] for row in matA]
print("The modified matrix is:\n",matM)

The same matrix is used here too. Another matrix called matM is created to store the new matrix. Each element in the matrix is iterated by element. If it is divisible by two, it is multiplied by three. Else, the same element is returned.

Modifying The Elements Using List Comprehension
Modifying The Elements Using List Comprehension

Adding Two Matrices

We can add the two given matrices and generate a new matrix using a for loop but again, it takes more number of lines.

matrixA = [[1,2,3],
           [4, 5, 6],
           [7, 8, 9]]
matrixB = [[6,5,4],
           [3,2,1],
           [9,8,7]]
res = []
for i in range(len(matrixA)):
    nrow = []
    for j in range(len(matrixA[i])):
        esum = matrixA[i][j] + matrixB[i][j]
        nrow.append(esum)
    res.append(nrow)
print("The new matrix is:")
print(res)

We are creating two matrices called matrixA and matrixB. The empty list called res is created to store the new matrix. We are going to go through every element in both matrices and add them. The result of this addition is stored in a variable called esum. We are also creating a new row to append the result of the addition. After every element of the first matrix is added to the corresponding element in the second matrix, the new row is appended to res matrix.

Adding Two Matrices Using For Loop
Adding Two Matrices Using For Loop

The same code using list comprehension is given below.

matrixA = [[1, 2, 3],
           [4, 5, 6],
           [7, 8, 9]]
matrixB = [[6,5,4],
           [3,2,1],
           [9,8,7]]
res = [[matrixA[i][j] + matrixB[i][j] for j in range(len(matrixA[i]))] for i in range(len(matrixB))]
print("The new matrix is:")
print(res)
Adding Two Matrices Using List Comprehension
Adding Two Matrices Using List Comprehension

Conclusion

A For Loop is an iterator used to iterate over a sequence of objects like strings, lists, tuples, and sets. It allows us to execute a set of statements over and over again for the items of the sequence. We have learned the different uses of for loops apart from iterating over a sequence of objects like appending multiple data frames into one and creating a nested dictionary.

For loop usually includes a lot of code lines to complete a task. There are other simpler ways to carry out the same task. One such method is list comprehension.

We looked at the possible ways of replacing a for loop with list comprehension in accessing and modifying the matrix elements and adding two matrices.

References

Learn more about list comprehension from the official documentation.

Find more about for loops here.