如何从3D Numpy数组便捷提取指定元素生成2D数组？

你可以用numpy的高级索引直接实现，完全不需要循环，代码简洁且效率更高：

import numpy as np

# 模拟你的数组A
A = np.random.rand(2, 512, 21128)
target_indices = [7, 41]

# 核心操作：按维度对应索引取值
result = A[np.arange(2), target_indices, :]

print(result.shape)  # 输出 (2, 21128)

原理说明 #

这里用到了numpy的广播式高级索引：

• np.arange(2)生成第一个维度的索引[0, 1]，对应数组的两个“行”
• target_indices是第二个维度的索引[7, 41]，分别对应每个“行”里要取的位置
• 两个索引数组会按位置一一匹配，取出A[0,7,:]和A[1,41,:]，并自动堆叠成形状为[2, 21128]的数组

这种矢量化操作比Python循环快得多，尤其是当数组规模较大时，性能优势会非常明显。

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