5 Best Ways to Find Diagonal Sum of a Matrix in Python

πŸ’‘ Problem Formulation: Calculating the diagonal sum of a matrix is a common task in linear algebra and programming. Given a square matrix as input, the desired output is the sum of all elements that form the main diagonal from the top-left to the bottom-right of the matrix. For example, in a 3×3 matrix with elements [1,2,3], [4,5,6], [7,8,9], the diagonal sum would be 1+5+9=15.

Method 1: Iterative Approach

This method involves iterating through the matrix with a loop to add up the elements on the main diagonal. It’s simple, straightforward, and the most intuitive approach to calculate the diagonal sum in Python. It works for any n x n square matrix.

Here’s an example:

matrix = [[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]]
          
diagonal_sum = 0
for i in range(len(matrix)):
    diagonal_sum += matrix[i][i]

Output:

15

The code snippet initializes a variable diagonal_sum to zero. Then, it uses a for loop to iterate through each row of the square matrix, adding the element at the current row and column index to the diagonal_sum. The loop index i ensures that we’re always referring to the proper diagonal element.

Method 2: Using the sum() Function with a Generator Expression

Python’s sum() function, paired with a generator expression, can calculate the diagonal sum in a compact, Pythonic way. It’s a concise alternative to the iterative approach and is suitable for those familiar with generator expressions.

Here’s an example:

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

diagonal_sum = sum(matrix[i][i] for i in range(len(matrix)))

Output:

15

The code snippet uses the sum() function to compute the sum of a generator expression. The expression iterates over the indexes of the matrix’s rows and calculates the appropriate diagonal element for each iteration, adding them up to find the total diagonal sum.

Method 3: Utilizing NumPy Library

For mathematic and scientific computations, the NumPy library comes in handy for working with arrays and matrices. NumPy’s trace() method provides a fast and efficient way to calculate the sum of the diagonals of an array.

Here’s an example:

import numpy as np

matrix = np.array([[1, 2, 3],
                   [4, 5, 6],
                   [7, 8, 9]])

diagonal_sum = np.trace(matrix)

Output:

15

This code snippet uses NumPy’s array() function to create a matrix and calls the trace() method to find the sum of the diagonals. It’s a clean, efficient, and concise way to calculate the diagonal sum, especially for larger matrices.

Method 4: Using the zip() Function

This method leverages Python’s built-in zip() function, which can be used to pair the i-th row and i-th column together, allowing us to extract diagonal elements and sum them.

Here’s an example:

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

diagonal_sum = sum(row[i] for i, row in enumerate(zip(*matrix)))

Output:

15

The code initiates the zip() function with unpacked rows of the matrix, which transposes the matrix. Then, it enumerates over the transposed matrix, selecting the diagonal elements using the current index i, and sums them up.

Bonus One-Liner Method 5: Using List Comprehension and enumerate()

Combining list comprehension with enumerate function, we can craft a one-liner that’s both readable and efficient to find the diagonal sum of a matrix.

Here’s an example:

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

diagonal_sum = sum(row[i] for i, row in enumerate(matrix))

Output:

15

The code snippet employs list comprehension to iterate over each row while using enumerate() to get both the index and the row itself. For each iteration, it takes the element at the position that matches the current row index and sums them up to find the diagonal sum.

Summary/Discussion

  • Method 1: Iterative Approach. It’s simple to understand and implement, making it great for beginners. However, it’s more verbose than other methods.
  • Method 2: Using sum() Function with Generator Expression. It’s more Pythonic and concise, but might be less readable to those unfamiliar with generator expressions.
  • Method 3: Utilizing NumPy Library. Ideal for large datasets and matrices, offering performance enhancements. Requires an external library.
  • Method 4: Using the zip() Function. A clever use of Python’s features, but may be less intuitive to understand at first glance.
  • Bonus Method 5: One-Liner Using List Comprehension and enumerate(). It’s extremely concise, but its readability depends on the user’s comfort with list comprehensions and enumerate().