按列索引的 array/list 访问数组元素

Question

我有二维整数数组（例如 A 与 A.shape (10,5)）和一维列索引数组（通常彼此不同）（例如 idx 与 idx.shape （ 10,)).从第 i 行开始，我想从数组 A 中获取列索引为 idx[i] 的元素。如果所需的输出是获取元素的一维数组或这些元素的列表，最好（最快）的解决方案是什么？

A = np.arange(50).reshape(10,5)
idx=np.array([2,4,0,0,3,1,3,1,2,4])

理想的输出：

output = [2,9,10,15,23,26,33,36,42,49]

或

output = np.array([2,9,10,15,23,26,33,36,42,49])

Answer 1

您可以使用 enumerate 访问您所在的行号，然后使用它访问您在 idx 上所需的索引，如下所示：

output = []

for row in enumerate(A):
    output.append(row[1][idx[row[0]]])

由此我得到了output == [2, 9, 10, 15, 23, 26, 33, 36, 42, 49]

Answer 2

使用numpy你可以获取列索引的对角线，例如：

np.diag(A[:,idx])
# array([ 2,  9, 10, 15, 23, 26, 33, 36, 42, 49])

Answer 3

A[np.arange(A.shape[0]), idx]

np.diag(A[:,idx]) 有效，但做了比必要更多的工作。 A[:,idx] 是 (10,10)array (in this example), which is then whittled down to a(10,)`.

对于这个小数组，我的速度快 2 倍。对于 (100,50) 它快 16 倍。

按列索引的 array/list 访问数组元素

Access array elements by array/list of column indexes

python

numpy