python - 创建一个数组，其中每个元素都存储其索引

我想创建一个2d numpy数组，其中每个元素都是其索引的元组。
示例（4x5）：

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

       [[1, 0],
        [1, 1],
        [1, 2],
        [1, 3],
        [1, 4]],

       [[2, 0],
        [2, 1],
        [2, 2],
        [2, 3],
        [2, 4]],

       [[3, 0],
        [3, 1],
        [3, 2],
        [3, 3],
        [3, 4]]])

我将创建一个pythonlist并理解以下列表：

[[(y,x) for x in range(width)] for y in range(height)]

有没有更快的方法来达到同样的效果呢？

最佳答案

下面是一个基于初始化的方法-

def create_grid(m,n):
    out = np.empty((m,n,2),dtype=int) #Improvement suggested by @AndrasDeak
    out[...,0] = np.arange(m)[:,None]
    out[...,1] = np.arange(n)
    return out

样本运行-

In [47]: create_grid(4,5)
Out[47]:
array([[[0, 0],
        [0, 1],
        [0, 2],
        [0, 3],
        [0, 4]],

       [[1, 0],
        [1, 1],
        [1, 2],
        [1, 3],
        [1, 4]],

       [[2, 0],
        [2, 1],
        [2, 2],
        [2, 3],
        [2, 4]],

       [[3, 0],
        [3, 1],
        [3, 2],
        [3, 3],
        [3, 4]]])

到目前为止发布在(4,5)网格和更大尺寸上的所有方法的运行时测试-

In [111]: %timeit np.moveaxis(np.indices((4,5)), 0, -1)
     ...: %timeit np.mgrid[:4, :5].swapaxes(2, 0).swapaxes(0, 1)
     ...: %timeit np.mgrid[:4,:5].transpose(1,2,0)
     ...: %timeit create_grid(4,5)
     ...:
100000 loops, best of 3: 11.1 µs per loop
100000 loops, best of 3: 17.1 µs per loop
100000 loops, best of 3: 17 µs per loop
100000 loops, best of 3: 2.51 µs per loop

In [113]: %timeit np.moveaxis(np.indices((400,500)), 0, -1)
     ...: %timeit np.mgrid[:400, :500].swapaxes(2, 0).swapaxes(0, 1)
     ...: %timeit np.mgrid[:400,:500].transpose(1,2,0)
     ...: %timeit create_grid(400,500)
     ...:
1000 loops, best of 3: 351 µs per loop
1000 loops, best of 3: 1.01 ms per loop
1000 loops, best of 3: 1.03 ms per loop
10000 loops, best of 3: 190 µs per loop