切片多维数组获得一个向量

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