Slicing A Multidimensional Array Without A For Loop

I have a 3D array r (1000 x 10 x 2000) constructed as follows:

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

This array, r, can be viewed as a 1000 x 10 matrix repeated 2000 times.

I would like to reduce this array according to the following rule:

• from each matrix select only the column which has the max value in the first row

The columns to be selected ca be obtained via: np.argmax(r[0], axis=0).

The result should be a 1000 x 2000 matrix.

I wonder if it is possible to get something like that without using a for loop or list comprehensions.

Here is a for loop which achieves the above task:

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

The solution I figured looks like this:

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

More elegant solutions are, of course, welcome.