如何使用numpy使用索引向量從矩陣中提取元素？

Question

假設我有一個mxn階的矩陣A和mx1的一個向量。我想通過使用向量的元素作爲每行中的偏移量來從矩陣A的每一行提取元素。

例如，

A = [[3, 0, 0, 8, 3], 
    [9, 3, 2, 2, 6], 
    [5, 5, 4, 2, 8], 
    [3, 8, 7, 1, 2], 
    [3, 9, 1, 5, 5]] 

和向量​​

y = [4,2,1,3,2] 

我想要實現的是，我想辦法來提取這樣的元素的向量索引的每個元素的元素在A的相應行中，即執行

for i in xrange(y): 
    A[i,y[i]] = #perform operations here 

不使用任何顯式循環。

的預期結果是，

[3,2,5,1,1] 

我使用Python以及numpy的圖書館。

感謝您的幫助！

Answer 1

您應該通過轉換列表A成NumPy的數組開始：

>>> import numpy as np 
>>> A = np.array([[3, 0, 0, 8, 3], 
...    [9, 3, 2, 2, 6], 
...    [5, 5, 4, 2, 8], 
...    [3, 8, 7, 1, 2], 
...    [3, 9, 1, 5, 5]]) 
... 
>>> y = [4, 2, 1, 3, 2] 

而在這之後，沒有什麼能阻止你使用advanced indexing：

>>> A[np.arange(A.shape[0]), y] 
array([3, 2, 5, 1, 1]) 
>>> A[np.arange(A.shape[0]), y] = -99 
>>> A 
array([[ 3, 0, 0, 8, -99], 
     [ 9, 3, -99, 2, 6], 
     [ 5, -99, 4, 2, 8], 
     [ 3, 8, 7, -99, 2], 
     [ 3, 9, -99, 5, 5]]) 

Answer 2

以下遵循傳統的python風格，但未針對速度進行優化。嘗試一下，看看它是否適用於你的用例。

A = [[1, 0, 0, 6, 1], 
    [8, 2, 1, 1, 5], 
    [1, 1, 0, 0, 4], 
    [3, 8, 7, 1, 2], 
    [3, 9, 1, 5, 5]] 

b = [0, 1, 2, 3, 4] 

result = [A[i][b_i] for i, b_i in enumerate(b)] 
# -> [1, 2, 0, 1, 5]