如何无需循环从数组X每行按索引数组I提取对应列的值

当然可以不用循环搞定！你之前用X[:,I]没得到预期结果，是因为这个索引方式的逻辑和你想要的不一样——它会把每行都按照I里的索引提取所有对应列，最后得到一个二维数组，而不是每行只取一个元素的一维结果。下面分两种常见场景给你解决方案：

场景1：使用NumPy数组（最常用的数值计算场景） #

NumPy的高级索引可以完美解决这个问题，核心是给每行匹配对应的行索引，再和列索引数组I配对：

1. 先生成和X行数一致的行索引数组：用np.arange(X.shape[0])，它会生成[0,1,2,...,n-1]（n是X的行数）
2. 将行索引数组和I数组一起传入X的索引，就能实现每行取对应列的元素

示例代码：

import numpy as np

# 示例输入数组
X = np.array([[0.3, 0.4, 0.5],
              [0.1, 0.7, 0.9],
              [0.3, 0.6, 0.9]])
I = np.array([0, 2, 1])

# 正确提取方式
Y = X[np.arange(X.shape[0]), I]
print(Y)  # 输出: [0.3 0.9 0.6]

为什么这个方法有效？因为NumPy会把两个索引数组中对应位置的元素配对（比如第0行配第0列，第1行配第2列，第2行配第1列），直接定位到每个目标元素，完全不需要手动写循环。

场景2：使用纯Python列表 #

如果是普通的Python列表，也可以用列表推导式实现（这是Python内置的高效语法，不算手动写循环）：

X = [[0.3, 0.4, 0.5], [0.1, 0.7, 0.9], [0.3, 0.6, 0.9]]
I = [0, 2, 1]

Y = [row[idx] for row, idx in zip(X, I)]
print(Y)  # 输出: [0.3, 0.9, 0.6]

补充：为什么X[:,I]不符合预期？ #

你之前尝试的X[:,I]，实际效果是把每行的第0、2、1列都提取出来，最终得到一个和原数组行数相同、列数和I长度相同的二维数组：

print(X[:,I])
# 输出:
# [[0.3 0.5 0.4]
#  [0.1 0.9 0.7]
#  [0.3 0.9 0.6]]

这显然和你“每行只取一个元素”的需求不符，所以需要用上面的高级索引或者列表推导式来实现。

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