[英]fastest way to create numpy 2d array of indices
我想創建一個包含單元格索引的numpy 2d數組,例如可以使用以下命令創建這樣的2x2 mat:
np.array([[[0,0],[0,1]],[[1,0],[1,1]]])
換句話說,索引i,j
處的單元格應包含列表[i,j]
。
我可以做一個嵌套循環來做到這一點,但我想知道是否有一個快速的pythonic方式來做到這一點?
對於使用NumPy的性能,我建議基於數組初始化的方法 -
def indices_array(n):
r = np.arange(n)
out = np.empty((n,n,2),dtype=int)
out[:,:,0] = r[:,None]
out[:,:,1] = r
return out
對於通用(m,n,2)
形狀的輸出,我們需要一些修改:
def indices_array_generic(m,n):
r0 = np.arange(m) # Or r0,r1 = np.ogrid[:m,:n], out[:,:,0] = r0
r1 = np.arange(n)
out = np.empty((m,n,2),dtype=int)
out[:,:,0] = r0[:,None]
out[:,:,1] = r1
return out
注意:另外,請閱讀本文后面的2019年附錄 。 大m
, n
增強。
樣品運行 -
In [145]: n = 3
In [146]: indices_array(n)
Out[146]:
array([[[0, 0],
[0, 1],
[0, 2]],
[[1, 0],
[1, 1],
[1, 2]],
[[2, 0],
[2, 1],
[2, 2]]])
如果您需要2
列2D
數組,只需重塑 -
In [147]: indices_array(n).reshape(-1,2)
Out[147]:
array([[0, 0],
[0, 1],
[0, 2],
[1, 0],
[1, 1],
[1, 2],
[2, 0],
[2, 1],
[2, 2]])
時間和驗證 -
In [141]: n = 100
...: out1 = np.array(list(product(range(n), repeat=2))).reshape(n,n,2)
...: out2 = indices_array(n)
...: print np.allclose(out1, out2)
...:
True
# @Ofek Ron's solution
In [26]: %timeit np.array(list(product(range(n), repeat=2))).reshape(n,n,2)
100 loops, best of 3: 2.69 ms per loop
In [27]: # @Brad Solomon's soln
...: def ndindex_app(n):
...: row, col = n,n
...: return np.array(list(np.ndindex((row, col)))).reshape(row, col, 2)
...:
# @Brad Solomon's soln
In [28]: %timeit ndindex_app(n)
100 loops, best of 3: 5.72 ms per loop
# Proposed earlier in this post
In [29]: %timeit indices_array(n)
100000 loops, best of 3: 12.1 µs per loop
In [30]: 2690/12.1
Out[30]: 222.31404958677686
200x+
為200x+
加速, n=100
,基於初始化!
我們也可以使用np.indices
-
def indices_array_generic_builtin(m,n):
return np.indices((m,n)).transpose(1,2,0)
計時 -
In [115]: %timeit indices_array_generic(1000,1000)
...: %timeit indices_array_generic_builtin(1000,1000)
100 loops, best of 3: 2.92 ms per loop
1000 loops, best of 3: 1.37 ms per loop
np.array(list(product(range(n), repeat=2))).reshape(n,n,2)
這很有效
你想要np.ndindex
。
def coords(row, col):
return np.array(list(np.ndindex((row, col)))).reshape(row, col, 2)
coords(3, 2)
Out[32]:
array([[[0, 0],
[0, 1]],
[[1, 0],
[1, 1]],
[[2, 0],
[2, 1]]])
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.