按數組/列索引訪問數組元素

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倍。