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

Question

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

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

使用numpy，您可以采用列索引的对角线，例如：

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

Answer 2

您可以使用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 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倍。