Обрезка многомерного массива numpy для получения вектора

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