python - 在 Python Numba/NumPy 中实现的分摊 O(1) 滚动最小值

Question

我正在尝试实现具有摊销 O(1) get_min() 的滚动最小值.摊销的 O(1) 算法来自 the accepted answer in this post

---

原始函数:

import pandas as pd
import numpy as np
from numba import njit, prange

def rolling_min_original(data, n):
    return pd.Series(data).rolling(n).min().to_numpy()

---

我尝试实现摊销 O(1) get_min()算法:(这个函数对非小n有不错的表现)

@njit
def rollin_min(data, n):
    """
        brief explanations:

        param: stk2: the stack2 in the algorithm, except here it only stores the min stack
        param: stk2_top: it starts at n-1, and drops gradually until it hits -1 then it comes backup to n-1
            if stk2_top= 0 in the current iteration(it will become -1 at the end):
                that means stk2_top is pointing at the bottom element in stk2, 
            after it drops to -1 from 0, in the next iteration, stk2 will  be reassigned to a new array data[i-n+1:i+1],
            because we need to include the current index. 

        at each iteration:
        if stk2_top <0: (i.e. we have 0 stuff in stk2(aka stk2_top <0)
            - copy the past n items(including the current one) to stk2, so that stk2 has n items now
            - pick the top min from stk2(stk2_top = n-1 momentarily)
            - move down the pointer by 1 after the operation(n-1 becomes n-2)

        else: (i.e. we have j(1<=j<= n-1) stuff in stk2)
            - pick the top min from stk2(stk2_top is j-1 momentarily)
            - move down the pointer by 1 after the operation(j-1 becomes j-2)

    """


    if n >1:  

        def min_getter_rev(arr1):
            arr = arr1[::-1]
            result = np.empty(len(arr), dtype = arr1.dtype)
            result[0]= local_min = arr[0]

            for i in range(1,len(arr)):
                if arr[i] < local_min:
                    local_min = arr[i]
                result[i] = local_min
            return result

        result_min= np.empty(len(data), dtype= data.dtype)
        for i in prange(n-1):
            result_min[i] =np.nan


        stk2 = min_getter_rev(data[:n])
        stk2_top = n-2#it is n-2 because the loop starts at n(not n-1)which is the second non nan term
        stk1_min = data[n-1]#stk1 needs to be the first item of the stk1
        result_min[n-1]= min(stk1_min, stk2[-1])    

        for i in range(n,len(data)):
            if stk2_top >= 0:
                if data[i] < stk1_min:
                    stk1_min= min(data[i], stk1_min)#the stk1 min
                result_min[i] = min(stk1_min, stk2[stk2_top])#min of the top element in stk2 and the current element

            else:
                stk2 = min_getter_rev(data[i-n+1:i+1])
                stk2_top= n-1
                stk1_min = data[i]
                result_min[i]= min(stk1_min, stk2[n-1])

            stk2_top -= 1

        return result_min   
    else:
        return data

---

n 时的天真实现很小:

@njit(parallel= True)
def rolling_min_smalln(data, n):
    result= np.empty(len(data), dtype= data.dtype)

    for i in prange(n-1):
        result[i]= np.nan

    for i in prange(n-1, len(data)):
        result[i]= data[i-n+1: i+1].min()

    return result

---

一些用于测试的小代码

def remove_nan(arr):
    return arr[~np.isnan(arr)]

if __name__ == '__main__':

    np.random.seed(0)
    data_size = 200000
    data = np.random.uniform(0,1000, size = data_size)+29000

    w_size = 37

    r_min_original= rolling_min_original(data, w_size)
    rmin1 = rollin_min(data, w_size)

    r_min_original = remove_nan(r_min_original)
    rmin1 = remove_nan(rmin1)

    print(np.array_equal(r_min_original,rmin1))

---

函数rollin_min()具有几乎恒定的运行时间和比 rolling_min_original() 更低的运行时间什么时候n很大，很好。但是在n时表现不佳很低(在我的电脑中大约为 n < 37，在此范围内 rollin_min() 很容易被天真的实现击败 rolling_min_smalln() )。

我正在努力寻找改进的方法rollin_min() ，但到目前为止我被卡住了，这就是我在这里寻求帮助的原因。

---

我的问题如下:

我正在实现的算法是滚动/滑动窗口最小/最大的最佳算法吗？

如果不是，最好/更好的算法是什么？如果可以，如何从算法的角度进一步改进功能？

除了算法本身，还有哪些方法可以进一步提升函数的性能rollin_min() ？

---

编辑:根据多个请求将我的最新答案移至答案部分

Answer 1

代码运行缓慢的主要原因可能是在 min_getter_rev 中分配了一个新数组。您应该在整个过程中重复使用相同的存储空间。

然后，因为你真的不必实现一个队列，你可以做更多的优化。例如，两个堆栈的大小最多(通常)为 n，因此您可以将它们保存在大小为 n 的同一个数组中。从头开始种植一个，从最后种植一个。

您会注意到有一个非常规则的模式 - 按顺序从头到尾填充数组，从末尾重新计算最小值，在重新填充数组时生成输出，重复...

这导致了一个实际上更简单的算法，其解释更简单，根本不涉及堆栈。这是一个实现，其中包含有关其工作原理的评论。请注意，我没有费心用 NaN 填充开头:

def rollin_min(data, n):

    #allocate the result.  Note the number valid windows is len(data)-(n-1)
    result = np.empty(len(data)-(n-1), data.dtype)

    #every nth position is a "mark"
    #every window therefore contains exactly 1 mark
    #the minimum in the window is the minimum of:
    #  the minimum from the window start to the following mark; and
    #  the minimum from the window end the the preceding (same) mark

    #calculate the minimum from every window start index to the next mark
    for mark in range(n-1, len(data), n):
        v = data[mark]
        if (mark < len(result)):
            result[mark] = v
        for i in range(mark-1, mark-n, -1):
            v = min(data[i],v)
            if (i < len(result)):
                result[i] = v

    #for each window, calculate the running total from the preceding mark
    # to its end.  The first window ends at the first mark
    #then combine it with the first distance to get the window minimum

    nextMarkPos = 0
    for i in range(0,len(result)):
        if i == nextMarkPos:
             v = data[i+n-1]
             nextMarkPos += n
        else:
            v = min(data[i+n-1],v)
        result[i] = min(result[i],v)

    return result