如何用每个条带的不同值初始化 numpy 二维数组

Question

我想初始化 2D numpy 数组，其中列小于行 + x（x 是参数）的条带为 2，列大于行 + y 的条带为 3（假设 y > x) 像这样：

有没有比暴力嵌套循环更快的方法：

pos_path = np.zeros((rows + 1, cols + 1), dtype=np.int32)
rows = 384;
cols = 288;
x = -12
y = 23
for r in range(0, rows+1):
    for c in range(0, cols + 1):
        if c < r + x:
             pos_path[r, c] = 2
        elif c > r + y:
             pos_path[r, c] = 3

Answer 1

我们可以简单地使用具有外部操作的范围数组来获得等效的掩码，然后使用它们以向量化的方式在输出中分配相关值，就像这样 -

pos_path_out = np.zeros((rows + 1, cols + 1), dtype=np.int32)
R = np.arange(0, rows+1)
C = np.arange(0, cols+1)
m1 = C < R[:,None] + x
m2 = C > R[:,None] + y
pos_path_out[m1] = 2
pos_path_out[m2] = 3

或者，我们可以直接将输出数组设置为关闭 m1 然后分配给 m2 -

pos_path_out = m1*2
pos_path_out[m2] = 3

或者一次性关闭m1和m2的最终输出-

pos_path_out = m1*2+m2*3

对于大型阵列，利用 numexpr -

的多核

import numexpr as ne

pos_path_out = ne.evaluate('(C < R2D + x)*2 + (C >R2D + y)*3',{'R2D':R[:,None]})

另一个利用蒙版的斜坡特性，我们可以创建具有范围外部比较的蒙版 -

pos_path_out = np.zeros((rows + 1, cols + 1), dtype=np.int32)
ra = np.arange(cols+1)
p1 = np.arange(y,rows+y+1)
p2 = np.arange(x,rows+x+1)
m1 = p1[:,None]<ra
m2 = p2[:,None]>ra
np.putmask(pos_path_out,m1,3)
np.putmask(pos_path_out,m2,2)

或使用np.tri创建这些面具-

m1 = ~np.tri(rows+1,cols+1,k=y, dtype=bool)
m2 = np.tri(rows+1,cols+1,k=x-1, dtype=bool)

如何用每个条带的不同值初始化 numpy 二维数组

How to initialize numpy 2D array with different values for each strip

python

numpy

vectorization

python-3.x