I have a 4x4 array, A1, and need to retrieve its diagonal elements without looping or calling np.diag(). What's a way of doing so? Appreciate your help!

A1 = np.array([ [1, 4, 6, 8],[2, 5, 7, 10],[3, 6, 9, 13], [11, 12, 16, 0]])

CodePudding user response：

by indexing using the indices of the diagonal, which are the indices of the non-zeros of an identity matrix.

import numpy as np

A1 = np.array([ [1, 4, 6, 8],[2, 5, 7, 10],[3, 6, 9, 13], [11, 12, 16, 0]])
diag_pos = np.eye(A1.shape[0],dtype=bool).nonzero()
print(A1[diag_pos])

[1 5 9 0]