Question

我想创建一个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]]])

我会使用以下列表理解创建一个python list：

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

是否有更快的方法来实现相同的目标，可能采用numpy方法？

Answer 1

你这样做是因为你需要它还是只是为了运动？在前一种情况下：

np.moveaxis(np.indices((4,5)), 0, -1)

np.indices正是其名称所暗示的。只有它以不同的方式排列轴。所以我们用moveaxis

移动它们

正如@Eric所指出的，这种方法的一个吸引人的特点是它可以在任意数量的维度上进行修改：

dims = tuple(np.multiply.reduceat(np.zeros(16,int)+2, np.r_[0, np.sort(np.random.choice(16, np.random.randint(10)))]))
# len(dims) == ?
np.moveaxis(np.indices(dims), 0, -1) # works

Answer 2

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

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

Answer 3

您可以为此目的滥用numpy.mgrid或meshgrid：

>>> import numpy as np
>>> np.mgrid[:4,:5].transpose(1,2,0)
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]]])

Answer 4

您可以使用numpy.mgrid并交换它的轴：

>>> # assuming a 3x3 array
>>> np.mgrid[:3, :3].swapaxes(-1, 0)
array([[[0, 0],
        [1, 0],
        [2, 0]],

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

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

这仍然与您想要的数组略有不同，因此您可以滚动轴：

>>> np.mgrid[:3, :3].swapaxes(2, 0).swapaxes(0, 1)
array([[[0, 0],
        [0, 1],
        [0, 2]],

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

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

鉴于有人对结果进行了计时，我还想提出一个基于numba的手动版本，“打败所有人”：

import numba as nb
import numpy as np

@nb.njit
def _indexarr(a, b, out):
    for i in range(a):
        for j in range(b):
            out[i, j, 0] = i
            out[i, j, 1] = j
    return out

def indexarr(a, b):
    arr = np.empty([a, b, 2], dtype=int)
    return _indexarr(a, b, arr)

定时：

a, b = 400, 500
indexarr(a, b)  # numba needs a warmup run
%timeit indexarr(a, b)                                  # 1000 loops, best of 3: 1.5 ms per loop
%timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1)  # 100 loops, best of 3: 7.17 ms per loop
%timeit np.mgrid[:a, :b].transpose(1,2,0)               # 100 loops, best of 3: 7.47 ms per loop
%timeit create_grid(a, b)                               # 100 loops, best of 3: 2.26 ms per loop

并在一个较小的阵列上：

a, b = 4, 5
indexarr(a, b)
%timeit indexarr(a, b)                                 # 100000 loops, best of 3: 13 µs per loop
%timeit np.mgrid[:a, :b].swapaxes(2, 0).swapaxes(0, 1) # 10000 loops, best of 3: 181 µs per loop
%timeit np.mgrid[:a, :b].transpose(1,2,0)              # 10000 loops, best of 3: 182 µs per loop
%timeit create_grid(a, b)                              # 10000 loops, best of 3: 32.3 µs per loop

正如所承诺的那样，它在性能方面“击败了所有人”： - ）

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

4 个答案: