带有索引数组的 numpy 3D 数组矢量化访问答案

【问题标题】：numpy 3D array vectorized access with arrays of indices带有索引数组的 numpy 3D 数组矢量化访问
【发布时间】：2020-11-22 17:05:46
【问题描述】：

我有一个 3D numpy 数组 A 代表一批图像： A.shape -> (batch_size, height, width)

我想使用另外两个大小为 batch_size 的数组 Hs、Ws 来访问这个数组。它们包含我要访问的每个图像的 x 索引和 y 索引。

示例 2 张 3x3 的图片：

A.shape(2,3,3)

A = [[[1,2,3],[5,6,7],[8,9,10]], [[10,20,30],[50,60,70],[ 80,90,100]]]

Hs = [0,2]

Ws = [1,2]

我想访问 A 以便得到：

A[:, Hs,Ws] = [2,100]

不幸的是，这样做 (A[:, Hs,Ws]) 会导致 2x2 数组 (batch_size x batch_size)

使用 for 循环执行，如下所示：

Result = np.zeros(batch_size)
for b in range(0,batch_size):
     Result[b] = A[b,Hs[b],Ws[b]]

是否可以通过以矢量化方式直接访问 A 而无需 for 循环来做到这一点？

【问题讨论】：

【解决方案1】：

你是这个意思吗：

In [6]: A = np.array(A); Hs=np.array(Hs); Ws=np.array(Ws)
In [7]: A.shape
Out[7]: (2, 3, 3)
In [8]: A[np.arange(2), Hs, Ws]
Out[8]: array([  2, 100])

使用索引数组时，它们会相互“广播”。这里有 (2,),(2,),(2,) 的广播很简单。

【讨论】：