我有一个维度为 (6,m,n)
的 3-D numpy 数组 a
。我还有一个 6 维 bool 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])]
有什么方法可以不使用循环来实现吗?
最佳答案
方法 #1
将 a
reshape 为 2D
,保持第一个轴长度不变。将每个这样的二维扁平 block 转换为元组,然后索引到 b
。此元组转换导致沿第一个轴将每个元素打包为一个索引器,以从 b
中选择一个元素。最后需要进行整形以获得 2D
输出。因此,实现看起来像这样 -
b[tuple(a.reshape(6,-1))].reshape(m,n)
或者,跳过所有 reshape 的困惑并简单地做 -
b[tuple(a)]
这会创建相同的索引器并解决问题。
方法 #2
或者,我们也可以计算扁平化索引,然后用这些索引到扁平化 b
中,并从中提取相关的 bool 值 -
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)
关于python - 使用索引数组中的索引元组索引多维数组 - NumPy/Python,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/56938555/