Pandas中基于Rolling窗口的复杂自定义函数实现求助

问题分析 #

你的核心需求是通过大小为3的滚动窗口处理DataFrame的a/b/c列，生成包含多维张量的新列array_col，但原代码存在几个关键问题：

• Pandas rolling.apply处理多列时，输入是展平的一维数组，而非二维数组
• 窗口数据需要逆序处理才能匹配你的示例要求
• 直接修改原数组会引发副作用，且numpy数组不支持列表的insert方法
• rolling.apply默认要求返回标量，需配置参数支持数组返回

修正后的完整代码 #

import pandas as pd
import numpy as np

# 定义张量积计算函数
def array_value(a, b, c):
    tem = np.tensordot(a, b, axes=0)
    tem = np.tensordot(tem, c, axes=0)
    return tem

# 定义滚动窗口处理逻辑
def make_coeff(x):
    # 将展平的窗口数据重塑为(3,3)的二维数组（3行数据，3列对应a/b/c）
    window_data = x.reshape(3, 3)
    
    # 提取窗口内的a/b/c列，并反转（匹配你示例中的窗口顺序：当前行→前一行→前两行）
    a_window = window_data[:, 0][::-1]
    b_window = window_data[:, 1][::-1]
    c_window = window_data[:, 2][::-1]
    
    # 处理a_new：在开头插入1
    a_new = np.insert(a_window, 0, 1)
    # 处理b_new：复制原数组后修改第一个元素为1
    b_new = b_window.copy()
    b_new[0] = 1
    # 处理c_new：复制原数组后修改第一个元素为1
    c_new = c_window.copy()
    c_new[0] = 1
    
    # 生成张量结果
    return array_value(a_new, b_new, c_new)

# 创建原始DataFrame
data = {'a': [1,2,3,4,5,6,7,8], 
        'b': [2,3,4,5,6,7,8,9], 
        'c': [3,4,5,6,7,8,9,10]}
df = pd.DataFrame(data)

# 生成array_col列：raw=True传入原始numpy数组，min_periods=3确保仅窗口满3条数据时计算
df['array_col'] = df[['a','b','c']].rolling(3, min_periods=3).apply(make_coeff, raw=True)

关键说明 #

1. 窗口数据处理：
   • 通过reshape(3,3)将展平的一维窗口数据还原为二维数组，方便提取a/b/c列
   • 使用[::-1]反转列数据，匹配你示例中窗口[3,2,1]的逆序要求
2. 数组操作修正：
   • 用np.insert替代列表insert方法处理numpy数组
   • 修改b_new/c_new前先复制原数组，避免修改原始窗口数据引发的副作用
3. Rolling参数配置：
   • raw=True：让rolling.apply直接传入numpy数组，提升处理效率并方便形状调整
   • min_periods=3：仅当窗口包含3条完整数据时才计算结果，前两行（索引0、1）会填充NaN

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