切片 NumPy ndarray 给出特定轴的索引

Question

Suppose there is a ndarray A = np.random.random([3, 5, 4]) , and I have another index ndarray of size 3 x 4, whose entry is the index I want to select from the 1st axis (the axis of dimension being 5). How can I achieve it using pythonic code?

Example:

A = [[[0.95220166 0.49801865 0.83217126 0.33361628]
      [0.31751156 0.85899736 0.81965214 0.62465746]
      [0.69251917 0.83201231 0.6089141  0.36589825]
      [0.96674647 0.6056233  0.45515703 0.90552863]
      [0.94524208 0.42422369 0.91633385 0.53177495]]

     [[0.02883774 0.18012477 0.64642352 0.21295456]
      [0.88475705 0.76020851 0.6888415  0.47958142]
      [0.17306953 0.94981064 0.91468365 0.37297622]
      [0.75924232 0.27537972 0.68803293 0.0904176 ]
      [0.14596762 0.70103752 0.06090593 0.07920207]]

     [[0.11092702 0.58002663 0.13553706 0.89662211]
      [0.09146413 0.86212582 0.65908978 0.2995175 ]
      [0.29025485 0.60788672 0.98595003 0.06762369]
      [0.56136928 0.09623415 0.20178919 0.46531331]
      [0.28628325 0.28215312 0.39670151 0.68243605]]]

Indices
  = [[3 1 2 1]
     [3 2 0 4]
     [3 3 1 2]]

Result_I_want 
       = [[0.96674647, 0.85899736, 0.6089141, 0.62465746]
          [0.75924232, 0.94981064, 0.64642352, 0.07920207]
          [0.56136928, 0.09623415, 0.65908978, 0.06762369]]

Answer 1

In [148]: A = np.arange(3*5*4).reshape([3, 5, 4])
In [151]: B = np.array([[3, 1, 2, 1],
     ...:      [3, 2, 0, 4],
     ...:      [3, 3, 1, 2]])
In [152]: B.shape
Out[152]: (3, 4)
In [153]: A.shape
Out[153]: (3, 5, 4)

Apply B to the middle dimension, and use arrays with shape (3,1) and (4,) for the other two. Together they broadcast to select a (3,4) array of elements.

In [154]: A[np.arange(3)[:,None],B,np.arange(4)]
Out[154]: 
array([[12,  5, 10,  7],
       [32, 29, 22, 39],
       [52, 53, 46, 51]])

Answer 2

Try np.take_along_axis :

A = np.arange(3*5*4).reshape([3, 5, 4])
# B is the same as your sample output
np.take_along_axis(A, B[:,None,:], axis=1).reshape(B.shape)

Output:

array([[12,  5, 10,  7],
       [32, 29, 22, 39],
       [52, 53, 46, 51]])