切片多維數組獲得一個向量

Question

在這個例子中，我試圖通過從多維數組中選擇相關元素來創建一個向量。

#data 
n=3 
rng = 4 
x = np.array([0,1,2],dtype=int) 
y = np.array([0,3,1],dtype=int) 
P = np.reshape(np.arange(n*rng*rng),(n,rng,rng)) 

output = np.zeros(n) 
for i in range(n): 
    output[i] = P[i,x[i],y[i]] 

這將返回

array([ 0., 23., 41.]) 

現在我想向量化上述操作。對我來說，合乎邏輯的做法是設置

output = P[0:n,x,y] 

但這返回

array([[ 0, 7, 9], 
     [16, 23, 25], 
     [32, 39, 41]]) 

任何人能解釋這是怎麼回事，我應該做些什麼來獲得預期的輸出？

在此先感謝

Answer 1

所有你需要的是：

>>> P[np.arange(n), x, y] 
array([ 0, 23, 41]) 

相關：Indexing Multi-dimensional arrays