获取 numpy ndarray 的切片(用于任意维度)

Question

我有一个任意维度的 Numpy 数组，以及一个索引向量，每个维度包含一个数字。我想获取对应于索引集的数组切片，小于所有维度的索引数组中的值，例如

A = np.array([[1, 2, 3, 4],
              [5, 6, 7, 8],
              [9,10,11,12]])
index = [2,3]
result = [[1,2,3],
          [5,6,7]]

直观的语法类似于 A[:index] ，但由于明显的原因，这不起作用。

如果数组的维度是固定的，我可以写A[:index[0],:index[1],...:index[n]] ;我可以使用某种列表理解，例如A[:i for i in index]？

Answer 1

您可以一次性对多个维度进行切片：

result = A[:2,:3]

将维度 1 切片到索引 2，将维度 2 切片到索引 3。

如果您有任意尺寸，则还可以创建切片tuple：

slicer = tuple(slice(0, i, 1) for i in index)
result = A[slicer]

切片定义start （0）、stop（您指定的索引）和step （1） - 基本上类似于range，但可用于索引。元组的第 i 个条目对数组的第 i 维进行切片。

如果只指定 stop -index，则可以使用以下简写：

slicer = tuple(slice(i) for i in index)

如果您知道维度的数量，我会推荐第一个选项，如果您不知道，我会推荐最后一个选项。