每行选择N个元素，无需循环

Question

给定形状为(D, N, Q)的numpy数组items和形状为(N, P)的另一个索引数组ids，如何使用索引nq_ids创建形状为(D, N, P)的新数组my_items，如下所示：

# 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]]

使用numpy魔术而不是使用任何显式循环？下面是一个简单的例子：

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)

如果可能的话，我还想保留元素顺序。你知道吗

Answer 1

现在应该可以工作了，返回与循环完全相同的ndarray：

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

然而，@hpaulj的方法速度更快，23.5µs比5µs快，所以使用它。你知道吗