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

Question

我有二维整数数组（例如A与A.shape（10,5））和1D列索引数组（通常彼此不同）（例如，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

使用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，）`。

对于这个小阵，我的速度要快两倍。对于(100,50)它快了16倍。