numpy 3D 数组矢量化访问与索引数组

Question

I have a 3D numpy array A representing a batch of images: A.shape -> (batch_size, height, width)

I want to access this array using two other arrays Hs,Ws, of size batch_size. They contain the x index and y index of each image that I want to access.

Example 2 images of size 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]

I want to acces A so that I get:

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

Doing it like this (A[:, Hs,Ws]) unfortunately results in a 2x2 array (batch_size x batch_size)

Executed with a for loop this would look like this:

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

Is it possible to do this without a for loop by accessing A directly in a vectorized manner?

Answer 1

Do you mean this:

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

When using indexing arrays, they 'broadcast' against each other. Here with (2,),(2,),(2,) the broadcasting is eash.