python - 获取 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]

将第一维切片到索引 2，将第二维切片到索引 3。

如果您有任意维度，您还可以创建切片的元组:

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)

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