如何对DataFrame的数组类型列使用pd.rolling计算逐元素滚动均值

实现思路 #

pandas的pd.rolling方法原生不支持元素为列表的列，我们可以先把data列的列表拆分为独立的数值列，完成滚动均值计算后再合并回列表格式即可。

完整实现代码 #

首先构造测试用的DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'time': range(10),
    'data': [[8,8,1], [3,3,4], [4,7,0], [1,4,4], [0,3,5], [1,1,6], [2,0,0], [0,9,3], [2,9,0], [2,8,0]]
})

然后拆分列+计算滚动均值+合并回列表：

# 1. 将data列的列表拆分为独立的数值列
expanded_data = pd.DataFrame(df['data'].tolist(), index=df.index)
# 2. 计算窗口为3的逐列滚动均值，round(1)保留1位小数和示例效果一致
rolling_res = expanded_data.rolling(3).mean().round(1)
# 3. 将每行的计算结果合并为列表，前两行缺失值设为NaN
df['rolling'] = rolling_res.apply(list, axis=1).mask(rolling_res.isna().all(axis=1))

输出验证 #

执行print(df.head())可以看到前几行输出符合预期：

time       data          rolling
0     0  [8, 8, 1]              NaN
1     1  [3, 3, 4]              NaN
2     2  [4, 7, 0]     [5.0, 6.0, 1.7]
3     3  [1, 4, 4]  [2.7, 4.7, 2.7]
4     4  [0, 3, 5]  [1.7, 4.7, 3.0]

如果需要自定义精度，调整round()的参数即可。

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