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

Question

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

我想使用第一个数组中每个位置 (m,n) 的 6 个值来检索第二个数组中的相应值。实际上，我会将 3D int 数组压缩为 2D boolean 数组。我认为解决方案是使用 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-flattened-block 转换为一个元组，然后索引到 b。 tuple-conversion 导致沿第一个轴将每个元素打包为 索引器 到 select 一个元素，每个元素离开 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

Indexing multi-dimensional array with tuple of indices from an indexing array - NumPy / Python

python

numpy

vectorization