Question

我有一个尺寸为a的3-D numpy数组(6,m,n)。我还有一个尺寸为b的6维布尔numpy数组(20,20,20,20,20,20)，可以有效地用作遮罩。

我想使用第一个数组中每个位置(m,n)的6个值来检索第二个数组中的对应值。有效地，我将3D int数组压缩为2D布尔数组。我以为解决方案将使用np.where，但我认为它不能使用值作为索引来处理。

简单的实现方式如下：

for i in range(m):
    for j in range(n):
         new_arr[i,j]=b[tuple(a[:,i,j])]

有没有不用循环就可以实现的方法？

Answer 1

方法1

将a重塑为2D，使第一轴长度相同。将每个2D扁平化的块转换为元组，然后索引为b。该元组转换导致每个元素作为 indexer 沿第一轴打包，以选择每个b元素。最后，需要重塑以获得2D的输出。因此，实现看起来像这样-

b[tuple(a.reshape(6,-1))].reshape(m,n)

或者，跳过所有重塑混乱的事情，简单地做-

b[tuple(a)]

这将执行相同的索引器创建并解决问题。

方法2

或者，我们也可以计算平化的索引，然后使用那些索引索引到平化的b并从中提取相关的布尔值-

b.ravel()[np.ravel_multi_index(a,b.shape)]

大型数据集上的时间-

In [89]: np.random.seed(0)
    ...: m,n = 500,500
    ...: b = np.random.rand(20,20,20,20,20,20)>0.5
    ...: a = np.random.randint(0,20,(6,m,n))

In [90]: %timeit b[tuple(a)]
14.6 ms ± 184 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [91]: %timeit b.ravel()[np.ravel_multi_index(a,b.shape)]
7.35 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

使用索引数组中的索引元组索引多维数组-NumPy / Python

1 个答案: