如何在Pandas中按产品ID分组实现动态窗口滚动均值计算

问题描述 #

现有如下结构的DataFrame：

需求：新增mean_sales列，计算规则为：若当前行freq值为x，则mean_sales是同product_id分组内当前行上下各x行的sales均值。要求先按product_id和year排序，再分组计算。

用户编写的伪代码执行时触发错误：TypeError: 'DataFrameGroupBy' object does not support item assignment，代码如下：

import pandas as pd

data = {
    "product_id": [1, 1, 2, 3],
    "year": [2023, 2022, 2023, 2022],
    "sales": [100, 200, 10, 600],
    "freq": [2, 3, 2, 3],
}

df = pd.DataFrame(data)
unique_frequencies = df["freq"].unique().tolist()

for frequency in unique_frequencies:
  df_tmp = df[[df["freq"] == frequency]].sort_values(by=["product_id","year"], ascending=True).groupby("product_id")
  df_tmp["mean_sales"] = df_tmp["sales"].rolling(window_size=frequency, closed="left").mean()

  df_results = pd.concat([df_results, df_tmp]) 

错误原因 #

代码中直接对DataFrameGroupBy对象进行列赋值操作，而GroupBy对象是分组后的聚合容器，并非可直接修改的DataFrame，不支持[]赋值语法。此外，用户的逻辑存在误区：freq是每行的独立值，不能按统一的freq值批量处理，需逐行计算动态窗口范围。

正确实现方案 #

步骤说明 #

1. 先按product_id和year对整个DataFrame排序，确保分组内的行顺序正确。
2. 按product_id分组，对每个分组逐行计算动态窗口的均值：
   • 对每行，根据自身freq值x，确定窗口范围为当前行上下各x行（超出分组边界时取边界值）。
   • 提取窗口内的sales值计算均值，赋值给mean_sales列。

完整代码 #

import pandas as pd

data = {
    "product_id": [1, 1, 2, 3],
    "year": [2023, 2022, 2023, 2022],
    "sales": [100, 200, 10, 600],
    "freq": [2, 3, 2, 3],
}

df = pd.DataFrame(data)

# 第一步：按product_id和year排序
df = df.sort_values(by=["product_id", "year"], ascending=True).reset_index(drop=True)

# 定义分组内的计算函数
def compute_dynamic_mean(group):
    group_len = len(group)
    mean_values = []
    # 重置分组内索引，确保从0开始计数
    group = group.reset_index(drop=True)
    for idx, row in group.iterrows():
        x = row["freq"]
        # 计算窗口的起始和结束索引（超出边界则取边界）
        start_idx = max(0, idx - x)
        end_idx = min(group_len - 1, idx + x)
        # 提取窗口内的sales并计算均值
        window_sales = group.loc[start_idx:end_idx, "sales"]
        mean_val = window_sales.mean()
        mean_values.append(mean_val)
    group["mean_sales"] = mean_values
    return group

# 分组应用计算函数
df = df.groupby("product_id", group_keys=False).apply(compute_dynamic_mean)

print(df)

执行结果 #

product_id  year  sales  freq  mean_sales
0           1  2022    200     3       150.0
1           1  2023    100     2       150.0
2           2  2023     10     2        10.0
3           3  2022    600     3       600.0

关键细节说明 #

• 若需求为排除当前行的上下x行均值，只需修改窗口提取逻辑为window_sales = group.loc[start_idx:end_idx, "sales"].drop(idx)。
• 使用group_keys=False避免分组键作为额外列添加到结果中。
• 重置分组内索引是为了确保行索引从0开始，方便计算窗口范围。

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