将多维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]])

有谁能解释一下这里发生了什么，以及我应该做些什么来获得期望的输出？谢谢。

- Forzaa

Answer 1

你所需要的是：

以下内容即可：

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

相关链接：多维数组索引

- Ashwini Chaudhary

---

非常感谢！我一整天都在苦苦挣扎。 - Forzaa