获取 numpy 数组切片的索引

Question

如何从二维 numpy 数组的切片中获取索引数组？ 例如：

array = np.arange(20).reshape(5,4)
section = array[1:,1:]

section.shape 是 (4, 3)。我想要一个维度为(2, 4, 3) 的数组indexes，第一个轴是section 中array 中每个值的索引。

索引应如下所示：

array([[[1, 1, 1],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]], 

       [[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]]])

需要索引以便我可以修改它，然后用array[indexes[0], indexes[1]] 索引array。

我认为这可以通过对 np.mgrid 进行切片来解决......但我不确定如何做到这一点。还是有更好的解决方案？

Answer 1

使用np.mgrid：

>>> i,j = array.shape
>>> np.mgrid[1:i,1:j]
array([[[1, 1, 1],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]],

       [[1, 2, 3],
        [1, 2, 3],
        [1, 2, 3],
        [1, 2, 3]]])