循环遍历 2D numpy 数组（例如创建一条线）

Question

我经常发现自己必须在二维数组中创建一条线（或某种其他形状）。换句话说，除了 y = mx + c 之外，数组的值在任何地方都为零。 （除此之外 - 这种方法的动机，而不是在一维数组中存储一条线，是我的工作通常需要二维傅里叶变换，因此我需要除线/形状/等之外的任何地方的零）。

我通常的做法如下：

array = numpy.zeros((height, width))
for i, line in enumerate(array):
    for j, pixel in enumerate(line):
         if j == m*i + c:
             array[i,j] = 1

这很好用，但它并没有让我觉得特别像pythonic，而且当数组变大时它往往会变得很慢。所以，我的问题是一个相当笼统的问题 - 有人知道这样做的更好方法吗？

提前致谢！

Answer 1

你可以在这里使用broadcasting 来摆脱那些嵌套循环 -

import numpy as np
out = (np.arange(height) == m*np.arange(width)[:,None]+c)+0.0

作为验证正确性的示例，使用这些参数 -

height = 10
width = 10
m = 0.5;
c = 6;

你会 -

In [306]: array
Out[306]: 
array([[ 0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

In [307]: out
Out[307]: 
array([[ 0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  1.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.]])

Answer 2

函数np.fromfunction 设计用于可以从索引构造数组的情况，例如这种情况。

在你的情况下，

np.fromfunction(lambda i, j: j == m*i+c, (height, width), dtype=np.float)

将等同于您的方法，但使用 numpy 的例程而不是 Python for 循环。

简短演示：

import numpy as np
height, width = 10,10
m, c = 2, 4

a = np.zeros((height, width))
for i, line in enumerate(a):
    for j, pixel in enumerate(a):
         if j == m*i + c:
             a[i,j] = 1

b = np.fromfunction(lambda i, j: j == m*i+c, (height, width), dtype=np.float)

np.all(a==b)
# True
b.astype(np.int) # as type added to reduce output (no need for all the periods)
#array([[0, 0, 0, 0, 1, 0, 0, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 1, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
#       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

编辑：即使这个答案被接受了，我想指出 @Divakar's answer 在我的机器上快了大约 10 倍。如果您正在寻找速度：使用该答案如果您的问题很容易实现矢量化，就像 Divakar 所展示的那样（并非每个fromfunction 调用都可以轻松矢量化）。我赞成它，因为它是解决这个问题的好方法。

Answer 3

使用np.put，但您需要创建特定索引的列表，您可以通过列表理解来完成：

>>> np.put(arr,[j for j in range(arr.shape[1]) for i in range(arr.shape[0]) if j == m*i + c],1)

演示：

>>> np.put(arr,[j for j in range(arr.shape[1]) for i in range(arr.shape[0]) if j == 3*i + 1],1)
>>> arr
array([[ 0.,  1.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.]])
>>> np.put(arr,[j for j in range(arr.shape[1]) for i in range(arr.shape[0]) if j == 0.5*i + 2],1)
>>> arr
array([[ 0.,  1.,  1.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.],
       [ 0.,  0.,  0.]])