获取一个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]?

- Mollusq

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 索引，你可以使用简写：

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

我建议如果您知道维度的数量，则选择第一个选项，否则选择最后一个选项。

- MSeifert