Python/Pandas：基于时变窗口的数组滚动求和实现问询

当然可以实现这种时变窗口的滚动求和需求！核心是针对数组的每个位置，使用对应的动态窗口长度来计算该位置的滚动和，下面我会一步步讲清楚实现思路和代码示例。

核心逻辑 #

对于数组arr的第i个元素（索引从0开始），对应的窗口长度是windows[i]：

• 我们需要计算arr中从max(0, i - windows[i] + 1)到i的所有元素之和
• 用max(0, ...)是为了避免窗口长度超过当前元素个数时出现索引越界的问题（比如第一个元素窗口长度为2的话，就只能取第一个元素本身求和）

基础实现（Python） #

先从最直观的版本开始，适合小规模数组：

def time_varying_rolling_sum(arr, windows):
    # 先校验输入：数组和窗口列表长度必须一致
    if len(arr) != len(windows):
        raise ValueError("arr 和 windows 的长度必须完全匹配")
    
    result = []
    for idx in range(len(arr)):
        window_size = windows[idx]
        # 计算求和的起始索引
        start = max(0, idx - window_size + 1)
        # 对切片求和并加入结果
        current_sum = sum(arr[start:idx+1])
        result.append(current_sum)
    
    return result

# 测试你给出的例子
# 场景1: t=2（对应索引1），arr=[1,2]，窗口值为2
arr_example1 = [1,2]
windows_example1 = [1,2]
print(time_varying_rolling_sum(arr_example1, windows_example1))  # 输出: [1, 3]

# 场景2: t=3（对应索引2），arr=[1,2,3]，窗口值为1
arr_example2 = [1,2,3]
windows_example2 = [1,1,1]
print(time_varying_rolling_sum(arr_example2, windows_example2))  # 输出: [1,2,3]
# 如果是前两个位置窗口为2，第三个为1：
windows_example3 = [2,2,1]
print(time_varying_rolling_sum(arr_example2, windows_example3))  # 输出: [1, 3, 3]

优化版本（适合大规模数组） #

如果你的数组很大，上面的sum循环会重复计算元素，效率较低。可以用前缀和数组来优化，把每个求和操作降到O(1)时间：

def optimized_time_varying_rolling_sum(arr, windows):
    if len(arr) != len(windows):
        raise ValueError("arr 和 windows 的长度必须完全匹配")
    
    # 构建前缀和数组：prefix_sum[0] = 0，prefix_sum[k] = arr[0]+arr[1]+...+arr[k-1]
    prefix_sum = [0] * (len(arr) + 1)
    for idx in range(len(arr)):
        prefix_sum[idx+1] = prefix_sum[idx] + arr[idx]
    
    result = []
    for idx in range(len(arr)):
        window_size = windows[idx]
        start = max(0, idx - window_size + 1)
        # 利用前缀和快速计算区间和：sum(arr[start:idx+1]) = prefix_sum[idx+1] - prefix_sum[start]
        current_sum = prefix_sum[idx+1] - prefix_sum[start]
        result.append(current_sum)
    
    return result

这个版本只需要遍历数组两次：一次构建前缀和，一次计算每个位置的滚动和，比基础版本的效率高很多，尤其当数组元素数量上万甚至更多时效果明显。

内容的提问来源于stack exchange，提问作者Michael