Python中提取多维数组内对角矩阵的对角线元素至单个数组

针对你给出的形状为(3, 4, 4)的3D数组，这里有几种简洁的NumPy方法可以提取每个子矩阵的对角线元素，得到目标数组：

方法1：使用np.diagonal()函数 #

这是最直接的方法，通过指定子矩阵对应的轴参数即可：

import numpy as np

A = np.array([[[1 , 0, 0, 0],
               [0 , 2, 0, 0],
               [0 , 0, 3, 0],
               [0 , 0, 0, 4]],

              [[5 , 0, 0, 0],
               [0 , 6, 0, 0],
               [0 , 0, 7, 0],
               [0 , 0, 0, 8]],

             [[9 , 0, 0, 0],
              [0 , 10, 0, 0],
              [0 , 0, 11, 0],
              [0 , 0, 0, 12]]])

# 提取每个子矩阵的对角线，axis1和axis2对应子矩阵的行、列轴
B = A.diagonal(axis1=1, axis2=2)
print(B)
# 输出：
# [[ 1  2  3  4]
#  [ 5  6  7  8]
#  [ 9 10 11 12]]

方法2：利用数组索引直接取值 #

对角线元素的位置是每个子矩阵中行索引等于列索引的位置，通过索引匹配即可：

# 生成0-3的索引序列，对应每个子矩阵的对角线位置
diag_indices = np.arange(4)
B = A[:, diag_indices, diag_indices]
print(B)
# 输出结果与方法1一致

方法3：使用np.take_along_axis()（灵活适配不同场景） #

如果需要更灵活地指定提取位置，可结合该函数：

# 创建与A的轴2匹配的索引数组
idx = np.arange(4)[None, :, None]
# 沿轴2提取对应索引的元素，再压缩维度
B = np.take_along_axis(A, idx, axis=2).squeeze(axis=2)
print(B)
# 输出结果与前两种方法一致

内容的提问来源于stack exchange，提问作者learning statistics