获取一个 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-indices 你可以使用简写：

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

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