if __name__=="__main__":
n = np.diag_indices(5)
print(n)
x = np.arange(0, 25).reshape((5, 5))
print(x)
print(x[n])

#output:
(array([0, 1, 2, 3, 4]), array([0, 1, 2, 3, 4]))
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]
[ 0  6 12 18 24]

此函数主要是返回一个对角线上的索引。接受两个参数，一个是n表示列表长度，一个是ndim，表示对角线索引的维度，默认为2.