Вызов функции с постоянным временем, повторяющейся в массиве Numpy, приводит к очень медленному коду

Question

У меня есть следующий фрагмент кода, который по существу выполняет следующие действия: учитывая двумерный массив, arr, вычисляет sum_arr следующим образом:

sum_arr[i, j] = arr[i, j] + min(sum_arr[i - 1, j-1:j+2]) if (i>0) else arr[i, j]

(разумные индексы для j - 1 : j + 2 изКонечно, все в пределах 0 и w)

Вот моя реализация:

import numpy as np

h, w = 1000, 1000 # Shape of the 2d array
arr = np.arange(h * w).reshape((h, w)) 

sum_arr = arr.copy()

def min_parent(i, j):
    min_index = j    
    if j > 0:
        if sum_arr[i - 1, j - 1] < sum_arr[i - 1, min_index]:
            min_index = j - 1
    if j < w - 1:
        if sum_arr[i - 1, j + 1] < sum_arr[i - 1, min_index]:
            min_index = j + 1    
    return (i - 1, min_index)


for i, j in np.ndindex((h - 1, w)):
    sum_arr[i + 1, j] += sum_arr[min_parent(i + 1, j)]

И вот проблема: этот фрагмент кода занимает слишком много времени для выполнения только для 1e6 операций (В среднем около 5 секунд на моей машине)

Как лучше это реализовать?

Answer 1

Использование динамического программирования: в другом массиве предварительно вычислите минуты для блоков размера X (в вашем случае вы делаете это для размера 3 (поскольку вы проверяете j-1, j, j + 1). Чтобы определитьмин для блока, используйте значение ссылочной позиции в исходном массиве и мин предыдущего блока, потому что вы, кажется, делаете это динамически.

Таким образом, вы просто назначаете индекс, который в этом нуждается.

Answer 2

Хотя ваша операция является последовательной в разных строках, в строках это не так.Поэтому легко векторизовать построчно и сохранить только 1D внешний цикл, который в относительном выражении не должен вызывать слишком много накладных расходов.

Действительно, это дает мне ~ 200x ускорение:

5.2975871179951355   # OP
0.023798351001460105 # vectorized rows

А код на самом деле довольно прост:

import numpy as np

h, w = 1000, 1000 # Shape of the 2d array
arr = np.arange(h * w).reshape((h, w)) 

def min_parent(i, j, sum_arr):
    min_index = j    
    if j > 0:
        if sum_arr[i - 1, j - 1] < sum_arr[i - 1, min_index]:
            min_index = j - 1
    if j < w - 1:
        if sum_arr[i - 1, j + 1] < sum_arr[i - 1, min_index]:
            min_index = j + 1    
    return (i - 1, min_index)

def OP():
    sum_arr = arr.copy()

    for i, j in np.ndindex((h - 1, w)):
        sum_arr[i + 1, j] += sum_arr[min_parent(i + 1, j, sum_arr)]
    return sum_arr

def vect_rows():
    h, w = arr.shape
    if w==1:
        return arr.cumsum(0)
    out = np.empty_like(arr)
    out[0] = arr[0]
    for i in range(1, h):
        out[i, :-1] = np.minimum(out[i-1, :-1], out[i-1, 1:])
        out[i, 1:] = np.minimum(out[i, :-1], out[i-1, 1:])
        out[i] += arr[i]
    return out

assert np.allclose(OP(), vect_rows())

from timeit import repeat

print(min(repeat(OP, number=3)))
print(min(repeat(vect_rows, number=3)))