Как инициализировать двумерный массив с разными значениями для каждой полосы

Question

Я хочу инициализировать двумерный массив с цифрами 2 для полосы, для которой столбец меньше строки + x (x является параметром), и с 3 для полосы, для которой столбец больше строки + y (предположим, 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]})

---

Еще один, использующий линейную природу масок, мы могли бы создать эти маски сexternal-сравнение -

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)