如何优化pandas滚动窗口中首个非零值与末尾值差值的计算效率

你原来的代码效率低的核心原因是rolling.apply默认走Python层循环，每个长度为3000的窗口都要单独执行一次lambda函数，20万行数据对应近20万次Python级调用，开销极高。下面给两个落地的优化方案：

方案1：NumPy向量化实现（无额外依赖，速度最快） #

利用二分查找快速定位每个窗口的首个非零值，所有逻辑都在C层执行，20万行数据可以做到毫秒级返回结果：

import numpy as np

window_size = 3000
dist_arr = df["dist"].to_numpy()
n = len(dist_arr)
max_diff = 0

if n >= window_size and (dist_arr != 0).any():
    # 提取所有非零值的索引（天然有序）
    non_zero_indexes = np.where(dist_arr != 0)[0]
    # 生成所有窗口的左边界
    left_bounds = np.arange(n - window_size + 1)
    # 二分查找每个左边界对应的第一个非零值的位置
    pos_in_non_zero = np.searchsorted(non_zero_indexes, left_bounds)
    # 过滤有效窗口：存在非零值且首个非零值在窗口范围内
    valid_mask = (pos_in_non_zero < len(non_zero_indexes)) & (non_zero_indexes[pos_in_non_zero] <= left_bounds + window_size - 1)
    # 批量计算差值
    end_values = dist_arr[left_bounds + window_size - 1]
    first_non_zero_values = dist_arr[non_zero_indexes[pos_in_non_zero]]
    diffs = np.where(valid_mask, end_values - first_non_zero_values, 0)
    max_diff = diffs.max()

方案2：numba JIT编译（改动最小） #

如果不想修改业务逻辑，只要给rolling.apply加numba引擎参数即可，性能可以提升一个数量级以上：
首先安装依赖：

pip install numba

修改后代码：

df["dist"].rolling(3000).apply(
    lambda x: x[-1] - x[x != 0][0] if (x != 0).any() else 0,
    engine="numba",
    raw=True # 直接传入numpy数组而非Series，进一步降低开销
).dropna().max()

注意第一次运行会有一次编译开销，后续重复调用速度会非常快。

内容的提问来源于stack exchange，提问作者Andrea Barnabò