Python Numpy-二維數組中的3維索引，無循環

Question

我想使用索引數組idx構造形狀為(n,p,q)的數組V1 ，將具有相同形狀的數組應用於形狀為(p,q)的數組V0 。 下面是使用循環構造它的方法。

for i in range(n):
    V1[i,:,:] = V0[idx[i,:,:],range(q)]

換句話說， idx[i,:,:]數組包含V0第一維元素的索引。 我將其與第二維的關聯索引（在range(q)捕獲range(q) ，以沿最終數組V1的第一維獲取對應的元素。

我的問題如下：是否有一種方法可以通過使用廣播/索引技術來構建不循環的V1 ？

謝謝。

Answer 1

怎么樣

V1 = V0[idx, range(q)] #?

例：

import numpy as np

# set up dummy data
n,p,q = 3,4,5
V1 = np.empty((n,p,q))
V0 = np.random.rand(p,q)
idx = np.random.randint(0,n,(n,p,q))

# original
V1_old = V1.copy()
for i in range(n):
    V1_old[i,:,:] = V0[idx[i,:,:],range(q)]

# new
V1_new = V0[idx, range(q)]

# test
print(np.array_equal(V1_old, V1_new)) # True