python - 如何使用索引向量对 3d ndarray 进行切片

Question

考虑 3d numpy 数组:

ax1, ax2, ax3 = 3, 3, 2
arr = np.asarray(range(ax1*ax2*ax3)).reshape([ax1, ax2, ax3])
arr: [[[0, 1], [2, 3], [4, 5]],
      [[6, 7], [8, 9], [10, 11]],
      [[12, 13], [14, 15], [16, 17]]]

和索引向量idx = [0, 1, 2]。

我想通过以下语句使用 idx 对数组 arr 进行切片:

res = [arr[i, :idx[i]+1] for i in range(ax1)]
res: [[[0, 1]],
      [[6, 7], [8, 9],
      [[12, 13], [14, 15], [16, 17]]]

但是这种切片看起来很复杂。

numpy是否支持这样不使用循环的操作？我正在寻找类似 arr[range(ax1), :idx+1] 的内容。

Answer 1

您的问题是结果值不是矩形:您无法将其正确表示为数组。

如果您愿意仅采用不同的格式值，则可以通过boolean mask获得您需要的内容。 :

>>> mask = np.tri(3, 3, dtype=bool)
>>> arr[mask]
array([[ 0,  1],
       [ 6,  7],
       [ 8,  9],
       [12, 13],
       [14, 15],
       [16, 17]])

原则是，对于 [0;2]^2 中的每对索引，您是否应该采用该对:

>>> np.tri(3, 3, dtype=bool)
array([[ True, False, False],
       [ True,  True, False],
       [ True,  True,  True]], dtype=bool)

这导致了极其简洁的结果:

>>> arr[np.tri(3, 3, dtype=bool)]
array([[ 0,  1],
       [ 6,  7],
       [ 8,  9],
       [12, 13],
       [14, 15],
       [16, 17]])