python - 调用通过 Numpy 数组迭代的恒定时间函数会导致代码非常慢

Question

我有以下代码片段，它基本上执行以下操作：给定一个 2d numpy 数组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 秒）

有什么更好的方法来实现这一点？

标签： pythonpython-3.xperformancenumpy

Answer 1

虽然您的操作是跨行顺序的，但在行内却不是。因此，很容易按行进行矢量化并只保留一个 1D 外循环，相对而言不应该产生太多开销。

事实上，这样做给了我约 200 倍的加速：

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)))