在没有for循环的情况下切片多维数组

Question

我有一个3D数组r (1000 x 10 x 2000)，其构造如下：

q = np.random.normal(size=(10,2000))
r = np.random.normal(loc=q, size=(1000,10,2000))

此数组r可以看作是重复2000次的1000 x 10矩阵。

我想根据以下规则缩小此数组：

• 从每个矩阵中仅选择第一行具有最大值的列

要选择的列可以通过np.argmax(r[0], axis=0)获得。

结果应为1000 x 2000矩阵。

我想知道是否有可能在不使用for循环或列表理解的情况下获得类似的东西。

---

这是一个for循环，可以完成上述任务：

x = []
for i, idx in enumerate(np.argmax(r[0], axis=0)):
    x.append(r[:,idx,i])
x = np.array(x).T

Answer 1

我想到的解决方案如下：

r[:, np.argmax(r[0],axis=0), np.arange(2000)]

当然，欢迎使用更优雅的解决方案。