如何在Python中按用户计算基于历史时间步的滚动均值？

数据准备 #

原始DataFrame定义如下：

import pandas as pd
import numpy as np

df = pd.DataFrame({'user': ['user_1', 'user_2', 'user_3','user_1', 'user_2',  'user_3',
                            'user_1', 'user_2',  'user_3'],
                   'values': [[1, 0, 2, 0], [1, 8, 0, 2],[6, 2, 0, 0],
                              [5, 0, 2, 2], [3, 8, 0, 0],[6, 0, 0, 2],
                             [3, 1, 1, 3], [2, 4, 1, 0],[4, 2, 0, 1]],
                   'time': [1, 1, 1, 2, 2, 2, 3, 3, 3]})

数据展示：

user       values  time
0  user_1  [1, 0, 2, 0]     1
1  user_2  [1, 8, 0, 2]     1
2  user_3  [6, 2, 0, 0]     1
3  user_1  [5, 0, 2, 2]     2
4  user_2  [3, 8, 0, 0]     2
5  user_3  [6, 0, 0, 2]     2
6  user_1  [3, 1, 1, 3]     3
7  user_2  [2, 4, 1, 0]     3
8  user_3  [4, 2, 0, 1]     3

需求说明 #

需要按用户计算累积滚动逐元素均值：

• time=1时，直接取当前values值；
• time=N（N>1）时，取time=N-1的计算结果与当前values的对应元素求均值，得到新的结果。

以user_1为例：

• time=1: [1, 0, 2, 0]
• time=2: ([1,0,2,0] + [5,0,2,2])/2 = [3, 0, 2, 1]
• time=3: ([3,0,2,1] + [3,1,1,3])/2 = [3, 0.5, 1.5, 2]

错误尝试分析 #

原代码通过groupby后用np.vstack(x).mean(0)计算的是所有时间步的总均值，而非需求中的累积滚动均值：

result = (df.groupby('user')['values']
         .agg(lambda x: np.vstack(x).mean(0).round(2))
       )

该代码输出的是每个用户所有values的逐元素总平均，不符合需求。

解决方案 #

方法1：循环实现（直观易懂） #

通过groupby结合自定义函数，对每个用户的values序列进行逐步计算：

def rolling_elementwise_mean(values_series):
    # 将列表转换为numpy数组，方便逐元素运算
    arr = np.array(values_series.tolist())
    result = np.zeros_like(arr, dtype=np.float64)
    # 第一个时间步直接取原始值
    result[0] = arr[0]
    # 从第二个时间步开始，计算前一步结果与当前值的均值
    for i in range(1, len(arr)):
        result[i] = (result[i-1] + arr[i]) / 2
    # 转换为保留两位小数的列表
    return [row.round(2).tolist() for row in result]

# 应用到每个用户组，并将结果拆分为对应行
df['rolling_mean'] = df.groupby('user')['values'].apply(rolling_elementwise_mean).explode()

print(df)

输出结果：

user       values  time      rolling_mean
0  user_1  [1, 0, 2, 0]     1    [1.0, 0.0, 2.0, 0.0]
1  user_2  [1, 8, 0, 2]     1    [1.0, 8.0, 0.0, 2.0]
2  user_3  [6, 2, 0, 0]     1    [6.0, 2.0, 0.0, 0.0]
3  user_1  [5, 0, 2, 2]     2    [3.0, 0.0, 2.0, 1.0]
4  user_2  [3, 8, 0, 0]     2    [2.0, 8.0, 0.0, 1.0]
5  user_3  [6, 0, 0, 2]     2    [6.0, 1.0, 0.0, 1.0]
6  user_1  [3, 1, 1, 3]     3  [3.0, 0.5, 1.5, 2.0]
7  user_2  [2, 4, 1, 0]     3  [2.0, 6.0, 0.5, 0.5]
8  user_3  [4, 2, 0, 1]     3  [5.0, 1.5, 0.0, 1.0]

方法2：加权累积平均（高效无循环） #

利用权重公式避免循环，适合大数据量场景：

def rolling_elementwise_mean(values_series):
    arr = np.array(values_series.tolist())
    # 计算权重：每一步的当前值权重为1/2^i，前序结果权重累积
    weights = np.power(0.5, np.arange(len(arr))[::-1])
    # 计算加权累积和，再除以权重和得到均值
    result = (arr * weights[:, None]).cumsum(axis=0) / weights.cumsum()[:, None]
    return [row.round(2).tolist() for row in result]

df['rolling_mean'] = df.groupby('user')['values'].apply(rolling_elementwise_mean).explode()

该方法的计算逻辑与循环法完全一致，但效率更高。

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