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

我有一个三维numpy数组

a

，尺寸

（6，m，n）

。我还有一个6-D布尔numpy数组

b

，其尺寸

（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])]

有没有办法在不使用循环的情况下实现这一点？

方法#1

将

a

重塑为

2D

，保持第一个轴的长度不变。将每个二维扁平块转换为元组，然后索引到

b

。这个元组转换导致每个元素沿着第一个轴打包，作为索引器，以选择每个元素关闭

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)

---

像往常一样聪明谢谢你！为什么

tuple（）？我假设它只适用于一维。@rboz22它也适用于三维数组，如
b[tuple（a）]
方法所示。元组转换后，在这些元组之间按元素进行打包以形成索引器，并在
b
中分别选择一个元素。这也是在文章中提到的。对不起，我在编辑之前发表了评论。对于n-D数组，
tuple（）
沿哪个维度将值转换为索引<代码>元组（a）

似乎产生非常大的值，而不是

0@rboz22
tuple（a）
只将
a
中的元素收集到tuple中。如果
a
中的元素在0中