Select N elements from each row without looping

Question

Given a numpy array items of of shape (D, N, Q) and another array of indices ids of shape (N, P) , how can I make a new array my_items of shape (D, N, P) , by using the indices nq_ids , like the following:

# How can these loops be avoided?
my_items = np.zeros((D, N, P))
for n in range(N):
    for p in range(P):
        my_items[:, n, p] = items[:, n, ids[n, p]]

with numpy magic instead of using any explicit loops? Here is a minimal example:

import numpy as np

D, N, Q, P = 2, 5, 4, 3  # Reduced problem dimensions.
items = 1.0 * np.arange(D * N * Q).reshape((D, N, Q))  # Example data
ids = np.arange(0, N * P).reshape(N, P) % Q  # Example ids

# How can these loops be avoided?
my_items = np.zeros((D, N, P))
for n in range(N):
    for p in range(P):
        my_items[:, n, p] = items[:, n, ids[n, p]]

# print('items', items)
# print('ids', ids)
# print('my_items', my_items)

I would also like to preserve the element order if possible.

Answer 1

This should work now, returning the exact same ndarray as your loop:

np.stack([np.take(items[:,i,:], ids[i, :], axis=1) 
          for i in range(ids.shape[0])], axis=2).transpose((0,2,1))

However, @hpaulj's method is faster, by 23.5 µs vs 5 µs. So use that.