[英]fastest way to create numpy 2d array of indices
I want to create a numpy 2d array containning the cells indices, for example such 2x2 mat can be created using : 我想创建一个包含单元格索引的numpy 2d数组,例如可以使用以下命令创建这样的2x2 mat:
np.array([[[0,0],[0,1]],[[1,0],[1,1]]])
In other words cell at index i,j
should contain the list [i,j]
. 换句话说,索引i,j
处的单元格应包含列表[i,j]
。
I could make a nested loop to do it c way but i am wondering if there is a fast pythonic way to do that? 我可以做一个嵌套循环来做到这一点,但我想知道是否有一个快速的pythonic方式来做到这一点?
For performance with NumPy, I would suggest an array initialization based approach - 对于使用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
For a generic (m,n,2)
shaped output, we need some modifications : 对于通用(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
Note: Also, read - 2019 addendum later on in this post for perf. 注意:另外,请阅读本文后面的2019年附录 。 boost with large m
, n
. 大m
, n
增强。
Sample run - 样品运行 -
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]]])
If you needed a 2
columns 2D
array, simply reshape - 如果您需要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]])
Timings and verification - 时间和验证 -
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+
speedup there for n=100
with the initialization based one! 200x+
为200x+
加速, n=100
,基于初始化!
We can also use np.indices
- 我们也可以使用np.indices
-
def indices_array_generic_builtin(m,n):
return np.indices((m,n)).transpose(1,2,0)
Timings - 计时 -
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)
这很有效
You want np.ndindex
. 你想要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.