如何对pandas DataFrame应用可变窗口自定义公式生成新列

注：你给出的初始化代码df = DataFrame(randn(10,2),columns=list('A'))存在笔误，randn(10,2)生成10行2列数组，但仅传入1个列名，实际运行会报错，测试时调整为10行1列即可。

固定窗口（前后各1个点）实现 #

你需要的B列计算逻辑是中心对齐的滑动窗口运算，无需逐行写循环，直接通过序列移位即可高效实现：

import pandas as pd
import numpy as np
from pandas import DataFrame

# 复现测试数据
np.random.seed(42)
df = DataFrame(np.random.randn(10, 1), columns=list('A'))

# 移位获取相邻位置的值
a_prev = df['A'].shift(1)   # 上一行值，对应A[i-1]
a_curr = df['A']            # 当前行值，对应A[i]
a_next = df['A'].shift(-1)  # 下一行值，对应A[i+1]

# 按规则计算B列，首尾行因凑不齐窗口返回NaN
df['B'] = a_curr ** 2 + 2 * a_prev * a_next

运行后B列的计算结果和你给出的规则完全一致：索引1处的值为A[1]^2 + 2*A[0]*A[2]，索引2处为A[2]^2 + 2*A[1]*A[3]，以此类推。

可变相邻点数的通用实现 #

如果需要把前后参与计算的相邻点数设为可变参数，可以封装通用滑动窗口计算函数，支持自定义窗口大小、自定义计算逻辑：

def calc_centered_window(df, col='A', k=1, custom_calc=None):
    window_length = 2 * k + 1
    # 默认计算逻辑：即你示例中的规则，k=1时与固定窗口实现完全一致
    if custom_calc is None:
        def default_calc(window_arr):
            center_pos = k
            center_val = window_arr[center_pos]
            cross_prod_sum = 0
            # 累加前后m位的乘积
            for m in range(1, k+1):
                cross_prod_sum += window_arr[center_pos - m] * window_arr[center_pos + m]
            return center_val ** 2 + 2 * cross_prod_sum
        custom_calc = default_calc
    
    # 调用pandas原生rolling接口，底层做了性能优化，比逐行循环快1~2个数量级
    return df[col].rolling(
        window=window_length,
        center=True,
        min_periods=window_length  # 必须凑齐所有窗口点才计算，避免首尾行值不准
    ).apply(custom_calc, raw=True)

函数参数说明：

• df: 原始DataFrame
• col: 参与计算的列名
• k: 当前行前后各取k个相邻点参与计算，总窗口长度为2k+1
• custom_calc: 自定义窗口计算函数，入参为长度2k+1的numpy数组，数组中心位置对应当前行值

使用示例 #

• 传入k=1，就是你最初的前后各1个点的计算逻辑：

df['B_k1'] = calc_centered_window(df, col='A', k=1)

• 如果需要前后各2个点参与计算，直接改k值即可，无需修改其他逻辑：

df['B_k2'] = calc_centered_window(df, col='A', k=2)

• 如果需要更换计算规则（比如卷积、加权求和、自相关计算等），直接传入自定义的custom_calc函数即可，无需调整外层窗口逻辑。

补充说明 #

• 该实现基于Pandas原生rolling接口，大数据量下性能远高于逐行iterrows/apply的写法
• 首尾行因凑不齐2k+1个窗口点默认返回NaN，如果需要填充首尾值，可以直接在结果上调用fillna()方法，支持前向填充、后向填充、固定值填充等模式
• 如果需要调整窗口权重，直接在自定义计算函数中修改权重系数即可。

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