python - Python Numpy - 没有循环的二维数组中的 3 维索引

Question

我想使用具有相同形状的索引数组构造一个形状的数组V1，应用于形状的数组。用循环构造它的方法如下。(n,p,q)idxV0(p,q)

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

换句话说，idx[i,:,:]数组包含 的第 1 维元素的索引V0。我将它与在 中捕获的第二维的关联索引一起应用range(q)，以获取沿最终数组的第一个维度的相应元素V1。

我的问题如下：有没有一种方法可以V1通过使用广播/索引技术来构建而不循环？

谢谢你。

标签： pythonnumpyindexingbroadcastingn-dimensional

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