Question

我有一个形状为[10，10，10]的数组A。

现在，我想使用另一个包含索引的[10，10，10，3]形状的数组B访问该数组。

作为输出，我想获得形状为[10，10，10]的数组C。因此，B中的每个索引都由A中的相应值“替换”。

不幸的是，尽管遍历索引数组并逐步为每个元素获取A的每个对应值，但我找不到合适的答案来解决该问题。我正在寻找一种更有效的解决方案。

非常感谢！

Answer 1

有两种方法可以做您想要的事情。第一个使用循环，第二个不使用循环。第二个速度快了大约10倍。

解决方案1 -循环

import numpy as np
a = np.random.normal(0,1,(10,10,10)) # original array

b = np.random.randint(0,10, (10,10,10,3)) # array of positions

c = np.zeros((10,10, 10)) # new array

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        for k in range(a.shape[2]):
            c[i,j,k]=a[tuple(b[i,j,k])]

这种方法在我的笔记本电脑上耗时约4ms

以上内容可以作为比较的基准。现在，使用数组切片和无循环完成相同的工作。

解决方案2 -无循环。更有效的数组切片

a_original_shape = a.shape
# reshape b to make it (10**3, 3) 
# b_flat[:,0] hold all the i coords, b_flat[:,0] holds j coords etc
b_flat = b.reshape( (np.product(a_original_shape),)+(3,) )

# slice out the values we want from a. This gives us a 1D array
c2_flat = a[[b_flat[:,i] for i in range(3)]]

# reshape it back to the original form. 
# All values will be the correct place in this new array
c2 = c2_flat.reshape(a_original_shape)

这种方法在我的笔记本电脑上花费大约0.5毫秒

您可以使用

检查这些方法是否具有相同的作用

np.all(c2==c)
True

这里的第二种解决方案大约花费第一种解决方案时间的10％

使用numpy中的另一个索引数组访问数组

1 个答案: