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上切片以提取横截面。

Answer 2

您可以这样做：

import numpy as np

arr = np.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]]])


first, second = zip(*enumerate([2, 1, 0]))

result = arr[first, second, :]
print(result)

输出

[[10 11 12 13 14]
 [25 26 27 28 29]
 [40 41 42 43 44]]