`numpy:diag_indices`

source: /dev/numpy/build-install/usr/lib/python3.14/site-packages/numpy/lib/_index_tricks_impl.py :926

Signature

def diag_indices ( n , ndim = 2 )

Summary

Return the indices to access the main diagonal of an array.

Extended Summary

This returns a tuple of indices that can be used to access the main diagonal of an array a with a.ndim >= 2 dimensions and shape (n, n, ..., n). For a.ndim = 2 this is the usual diagonal, for a.ndim > 2 this is the set of indices to access a[i, i, ..., i] for i = [0..n-1].

Parameters

n : int: The size, along each dimension, of the arrays for which the returned indices can be used.
ndim : int, optional: The number of dimensions.

Examples

import numpy as np

✓

Create a set of indices to access the diagonal of a (4, 4) array:

di = np.diag_indices(4)
di
a = np.arange(16).reshape(4, 4)
a
a[di] = 100
a

✓

Now, we create indices to manipulate a 3-D array:

d3 = np.diag_indices(2, 3)
d3

✓

And use it to set the diagonal of an array of zeros to 1:

a = np.zeros((2, 2, 2), dtype=int)
a[d3] = 1

✓

✗

Aliases