【468】numpy 数组获取对角线数值

数组的索引也可以通过数组来添加

>>> a = np.arange(0, 25)
>>> a.reshape(5, 5)
array([[ 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]])
>>> a = a.reshape(5, 5)	      

# 5*5矩阵
>>> a      
array([[ 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]])

# 对角线 左上-右下
>>> a[np.arange(5), np.arange(5)]
array([ 0,  6, 12, 18, 24])

# 行
>>> a[1, np.arange(5)]  
array([5, 6, 7, 8, 9])
>>> a[1]
array([5, 6, 7, 8, 9])
>>> a[1, :]
array([5, 6, 7, 8, 9])

# 列
>>> a[:, 1]
array([ 1,  6, 11, 16, 21])
>>> a[np.arange(5), 1]
array([ 1,  6, 11, 16, 21])

# 对角线 右上-左下
>>> a[np.arange(5), 4 - np.arange(5)]	      
array([ 4,  8, 12, 16, 20])