Numpy三维数组按子矩阵关联行索引提取指定列元素的优化方法

你有一个形状为(3,2,2)的三维NumPy数组，想要从每个二维子矩阵的第0列中，分别提取第0行、第1行、第0行的元素，最终得到array([0,3,3])。但直接使用arr[:,[0,1,0],0]会触发广播机制，得到一个(3,3)的二维数组；虽然arr[range(arr.shape[0]),[0,1,0],0]能得到正确结果，但希望有更简洁的实现方式。

方法1：利用原广播结果的对角线提取 #

直接对arr[:,[0,1,0],0]的结果取对角线，即可得到目标元素：

import numpy as np

arr = np.array([[[0, 0],
                 [1, 1]],

                [[2, 0],
                 [3, 1]],

                [[3, 0],
                 [4, 1]]])

result = np.diag(arr[:, [0,1,0], 0])
# 输出：array([0, 3, 3])

原理：你原本的索引操作会生成一个(3,3)数组，其中第i行第i列的元素正好是第i个二维子矩阵中需要提取的值，取对角线即可直接筛选出这些元素。

方法2：使用np.take_along_axis（推荐） #

这是NumPy专门用于沿指定轴按索引提取元素的函数，语义更清晰：

# 先提取所有子矩阵的第0列，得到形状为(3,2)的数组
col0 = arr[..., 0]
# 构造对应每个子矩阵的行索引，形状为(3,1)以匹配轴方向
row_indices = np.array([[0], [1], [0]])
# 沿轴1提取对应索引的元素，再展平为一维数组
result = np.take_along_axis(col0, row_indices, axis=1).ravel()
# 输出：array([0, 3, 3])

如果要更紧凑，可以写成一行：

result = np.take_along_axis(arr[..., 0], np.array([[0],[1],[0]]), axis=1).flatten()

方法3：简化版高级索引 #

用np.arange替代range，写法更简洁，同时保持高级索引的逻辑：

result = arr[np.arange(arr.shape[0]), [0,1,0], 0]
# 输出：array([0, 3, 3])

np.arange(arr.shape[0])会自动生成与第一维度长度匹配的索引数组，比手动写range更符合NumPy的惯用风格。

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