Question

我想使用索引数组V1构建具有形状(n,p,q)的数组idx，该数组具有相同的形状，并应用于数组V0形状为(p,q)。用循环构造它的方法如下。

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

换句话说，idx[i,:,:]数组包含V0的第一维元素的索引。我将其与在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

Python Numpy-二维数组中的3维索引，无循环

1 个答案: