python - Numpy 3D 数组数据沿指定轴切片

Question

假设我有一个形状为 (3, 4, 5) 的数组，并希望使用索引数组 [2, 1, 0] 沿第二个轴进行切片。

我无法用文字解释我想做什么，所以请引用下面的代码和图:

>>> src = np.arange(3*4*5).reshape(3,4,5)
>>> index = [2,1,0]

>>> src
>>> array([[[ 0,  1,  2,  3,  4],
    [ 5,  6,  7,  8,  9],
    [10, 11, 12, 13, 14],
    [15, 16, 17, 18, 19]],

   [[20, 21, 22, 23, 24],
    [25, 26, 27, 28, 29],
    [30, 31, 32, 33, 34],
    [35, 36, 37, 38, 39]],

   [[40, 41, 42, 43, 44],
    [45, 46, 47, 48, 49],
    [50, 51, 52, 53, 54],
    [55, 56, 57, 58, 59]]])
>>> # what I need is:
    array([[[10, 11, 12, 13, 14]],  # slice the 2nd row (index[0])
           [[25, 26, 27, 28, 29]],  # 1st row (index[1])
           [[40, 41, 42, 43, 44]]])  # 0th row (index[2])

Answer 1

src[np.arange(src.shape[0]), [2, 1, 0]]
# src[np.arange(src.shape[0]), [2, 1, 0], :]
array([[10, 11, 12, 13, 14],
       [25, 26, 27, 28, 29],
       [40, 41, 42, 43, 44]])

---

我们需要计算 axis=0 的索引:

>>> np.arange(src.shape[0])
array([0, 1, 2])

我们已经有了 axes=1 的索引。然后我们切过 axis=3 以提取横截面。