Extract Diagonal Elements From Numpy Array

The Numpy library in Python comes with a number of useful functions to work with and manipulate the data in arrays. In this tutorial, we will look at how to extract the diagonal elements from a 2d array in Numpy.

How to get the diagonal elements in Numpy?

You can use the numpy built-in numpy.diag() function to extract the diagonal elements of a 2d Numpy array. Pass the array as an argument to the function.

The following is the syntax –

numpy.diag(v, k)

The numpy.diag() function takes the following parameters –

1. v – The 2d array to extract the diagonal elements from.
2. k – The diagonal to extract the elements from. It is 0 (the main diagonal) by default. Diagonals below the main diagonal have k < 0 and the ones above the main diagonal have k > 0.

It returns the extracted elements from the diagonal as a numpy array.

Examples

Let’s now look at examples of using the above syntax to get the diagonal elements of a 2d array.

First, we will create a Numpy array that we will use throughout this tutorial.

import numpy as np

# create a 2D numpy array
arr = np.array([
    [1, 2, 3],
    [4, 5, 6],
    [7, 8, 9],
    [10, 11, 12]
])
# display the matrix
print(arr)

Output:

[[ 1  2  3]
 [ 4  5  6]
 [ 7  8  9]
 [10 11 12]]

Here, we used the numpy.array() function to create a 2d array of shape 4×3 (having 4 rows and 3 columns).

Example 1 – Extract the elements on the default diagonal

Let’s now use the numpy.diag() function to get the diagonal elements for the 2d array created above. We will use the default diagonal (k = 0).

# get the diagonal elements
res = np.diag(arr)
# display the diagonal elements
print(res)

Output:

[1 5 9]

We get the diagonal elements of the passed array as a 1d numpy array. You can see that the returned array has the same values as the main diagonal.

Example 2 – Extract the elements on a custom diagonal

In the above example, we extracted the elements of the main diagonal.

The numpy.diag() function comes with an optional parameter, k that you can use to specify the diagonal you want to extract the elements from.

The below image better illustrates the different values of k (representing different diagonals) for our input array.

k is 0 by default. The diagonals below the main diagonal have k < 0 and the diagonals above it have k > 0.

Let’s now use the numpy.diag() function to get the elements on diagonal, k = -1.

# get the elements of the diagonal -1
res = np.diag(arr, k=-1)
# display the diagonal elements
print(res)

Output:

[ 4  8 12]

We get the elements for the k = -1 diagonal.

Alternative usage of the numpy.diag() function

In the above examples, we used the numpy.diag() function to extract the diagonal elements from a 2d array. You can also use the numpy.diag() function to create a diagonal matrix.

For example, if you pass a 1d array to the numpy.diag() function, it will return a 2d array with the passed array’s elements on the kth diagonal.

# create a 1d array of diagonal elements
ar = np.array([1, 2, 3])
# create a diagonal matrix
res = np.diag(ar)
# display the returned matrix
print(res)

Output:

[[1 0 0]
 [0 2 0]
 [0 0 3]]

We get a diagonal matrix from the 1d array.

Summary

In this tutorial, we looked at how to extract the diagonal elements of a 2d array in Numpy. The following are the key takeaways from this tutorial.

• Use the numpy.diag() function to get the diagonal elements of a 2d array.
• You can specify the diagonal for which you want the extract the elements using the optional parameter k. By default, it represents the main diagonal, k = 0.

You might also be interested in –

• Numpy – Get the Lower Triangular Matrix (With Examples)
• Get the First N Rows of a 2D Numpy Array
• Numpy – Remove Duplicates From Array