創建一個數組，其中每個元素存儲其索引

Question

我想創建一個2d numpy數組，其中每個元素都是其索引的元組。

例（4×5）：

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]]]) 

我會創建一個包含以下列表理解的蟒蛇list：

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

有一個更快的方法來達到同樣的，可能與numpy的方法呢？

Answer 1

這裏有一個初始化爲基礎的方法 -

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]]]) 

運行ŧ

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 2

可以濫用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 3

您可以使用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 

正如承諾它在性能方面「擊敗所有」:-)

Answer 4

你這樣做是因爲你需要它還是僅僅爲了運動？在前一種情況下：

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