切片 NumPy ndarray 在特定轴上给出索引答案

【问题标题】：Slicing NumPy ndarray giving indices at specific axis切片 NumPy ndarray 在特定轴上给出索引
【发布时间】：2021-12-15 17:44:06
【问题描述】：

假设有一个ndarray A = np.random.random([3, 5, 4])，并且我有另一个大小为 3 x 4 的索引ndarray，其条目是我要从第一个轴（维度轴为 5）中选择的索引。如何使用 pythonic 代码实现它？

例子：

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

【问题讨论】：

@DaniMesejo，感谢您的指示。我只是添加了一个示例。

标签： python numpy slice numpy-ndarray

【解决方案1】：

试试np.take_along_axis:

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

输出：

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

【讨论】：

【解决方案2】：

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)

将B 应用于中间维度，并将形状为 (3,1) 和 (4,) 的数组用于其他两个维度。他们一起broadcast 选择一个 (3,4) 元素数组。

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

【讨论】：