如何用Numpy方法按索引列表过滤嵌套数组？

你可以直接利用Numpy的高级索引完成需求，完全不需要循环，代码更简洁且效率更高。

核心写法是通过行索引数组与目标列索引列表配对，精准定位每个要提取的元素：

import numpy as np

x = np.array(
    [
        [1, 2],
        [3, 4]
    ]
)
i = [0, 1]

# 生成行索引序列，和i中的列索引一一对应
result = x[np.arange(x.shape[0]), i]
print(result)  # 输出: [1 4]

原理说明 #

np.arange(x.shape[0])会生成和子数组数量一致的行索引数组[0, 1]，它和你的列索引列表i = [0, 1]组合后，会分别取x[0, 0]和x[1, 1]两个元素，最终拼接成你需要的结果数组。

如果i本身已经是Numpy数组，用法完全相同，无需额外转换。

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