以数组为索引的 Numpy 切片答案

【问题标题】：Numpy slice with array as index以数组为索引的 Numpy 切片
【发布时间】：2012-11-07 15:48:36
【问题描述】：

我正在尝试将完整的索引集提取到一个 N 维立方体中，似乎 np.mgrid 正是我所需要的。例如，np.mgrid[0:4,0:4] 生成一个 4 x 4 矩阵，其中包含相同形状数组中的所有索引。

问题是我想根据另一个数组的形状在任意数量的维度上执行此操作。 IE。如果我有一个任意维度的数组a，我想做类似idx = np.mgrid[0:a.shape] 的操作，但不允许使用这种语法。

是否可以构造我需要的切片以使np.mgrid 工作？或者是否有其他一些优雅的方式来做到这一点？以下表达式可以满足我的需要，但它相当复杂，可能效率不高：

np.reshape(np.array(list(np.ndindex(a.shape))),list(a.shape)+[len(a.shape)])

【问题讨论】：

标签： python numpy

【解决方案1】：

我一般用np.indices:

>>> a = np.arange(2*3).reshape(2,3)
>>> np.mgrid[:2, :3]
array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])
>>> np.indices(a.shape)
array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])
>>> a = np.arange(2*3*5).reshape(2,3,5)
>>> (np.mgrid[:2, :3, :5] == np.indices(a.shape)).all()
True

【讨论】：

谢谢，这比mgrid 解决方案还要好。

【解决方案2】：

我相信以下内容可以满足您的要求：

>>> a = np.random.random((1, 2, 3))
>>> np.mgrid[map(slice, a.shape)]
array([[[[0, 0, 0],
         [0, 0, 0]]],


       [[[0, 0, 0],
         [1, 1, 1]]],


       [[[0, 1, 2],
         [0, 1, 2]]]])

它产生的结果与np.mgrid[0:1,0:2,0:3] 完全相同，只是它使用a 的形状而不是硬编码的尺寸。

【讨论】：