创建numpy二维索引数组的最快方法

Question

我希望创建一个包含单元格索引的numpy 2d数组，例如可以使用以下方式创建这样的2x2矩阵：

np.array([[[0,0],[0,1]],[[1,0],[1,1]]])

换句话说，索引为i,j的单元格应包含列表[i,j]。我可以用C方式创建嵌套循环来完成此操作，但我想知道是否有一种更快速、更Pythonic的方法？

- Ofek Ron

Answer 1

为了在 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+ 的速度提升是基于初始化的方法，当 n=100 时。

---

2019补充说明

我们也可以使用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

- Divakar

---

它比np.ndindex更快吗？ - Ofek Ron

@OfekRon，为所有方法添加了时间。看看吧！ - Divakar

如果请求的矩阵不是正方形怎么办？例如，如果您在ndindex中给出了行数和列数计数。 - Ofek Ron

def create_indices(rows,cols): r = np.arange(rows) c = np.arange(cols) out = np.empty((rows,cols,2),dtype=int) out[:,:,0] = r[:,None] out[:,:,1] = c print out return out - Ofek Ron

@OfekRon 请查看刚刚添加的 indices_array_generic。 - Divakar

Answer 2

np.array(list(product(range(n), repeat=2))).reshape(n,n,2)

这个有效

- Ofek Ron

---

单元格索引是列表，我需要一个包含每个单元格中的单元格索引的二维矩阵...我同意这会变成三维，但这正是我所需要的。 - Ofek Ron

Answer 3

你想要 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]]])

- Brad Solomon