求助：使用Python Pandas实现基于动态窗口的移动平均计算

我现在需要实现一个动态移动平均的计算逻辑，但写的代码结果不对，想请大家帮忙看看问题出在哪，以及怎么修正。

需求说明 #

我有一个包含a、b、c三列的DataFrame，要新增一列d，d的计算规则如下：

• 如果b的值是0，d为NaN；
• 如果b的值为正：取当前行开始往后b行的c值计算平均，但如果剩余行数不足b行（包括当前行），则d直接等于当前行的c值；
• 如果b的值为负：取当前行开始往前abs(b)行的c值计算平均，但如果前面行数不足abs(b)行，则d直接等于当前行的c值；

测试数据与预期结果 #

我的测试DataFrame如下：

import numpy as np
import pandas as pd

data = {
    'a': [1, 2, 3, 4, 5, 6],
    'b': [4, -2, 2, 0, 1, 4],
    'c': [100, 53, 48, 28, 18, 202]
}
df = pd.DataFrame(data)

预期的d列结果：

• index0：b=4，取c[0:4]（[100,53,48,28]）的平均，结果为57.25；
• index1：b=-2，取c[0:2]（[53,100]）的平均，结果为76.5；
• index2：b=2，取c[2:4]（[48,28]）的平均，结果为38；
• index3：b=0，d为NaN；
• index4：b=1，刚好取当前行c值18，结果为18；
• index5：b=4，后面无足够行，直接取当前c值202；

我的错误代码 #

但我写的代码得到的结果不符合预期，代码如下：

# Function to calculate the custom rolling average
def custom_rolling_average(window_values):
    if len(window_values) > 0:
        return window_values.mean()
    else:
        return np.nan

# Calculate the rolling average
df['d'] = np.nan
for i, row in df.iterrows():
    window_size = abs(row['b'])
    window_start = max(i - window_size + 1, 0)
    window_end = i
    window_values = df.loc[window_start:window_end, 'c'].astype(float)
    if window_size == 0:
        df.at[i, 'd'] = np.nan
    else:
        df.at[i, 'd'] = custom_rolling_average(window_values)

问题分析与修正方案 #

你的现有代码有几个关键问题导致结果不符：

1. 窗口方向处理错误：不管b是正还是负，都只往前取窗口（从window_start到i），但需求里b为正时应该往后取，b为负时才往前取；
2. 边界判断缺失：没有处理“窗口大小超过可用行数”的情况，比如当b为正但当前行后面的行数不足窗口大小时，应该直接返回当前行的c值；
3. 窗口范围计算错误：比如b=4（index0）时，现有代码只取了当前行的c值，完全不符合取4行的需求。

下面是修正后的代码，完全符合需求：

import numpy as np
import pandas as pd

data = {
    'a': [1, 2, 3, 4, 5, 6],
    'b': [4, -2, 2, 0, 1, 4],
    'c': [100, 53, 48, 28, 18, 202]
}
df = pd.DataFrame(data)

def calculate_d(row, df):
    b_val = row['b']
    c_val = row['c']
    idx = row.name  # 获取当前行的索引
    
    if b_val == 0:
        return np.nan
    
    window_size = abs(b_val)
    total_rows = len(df)
    
    if b_val > 0:
        # 正向窗口：从当前索引开始往后window_size行
        end_idx = idx + window_size
        if end_idx > total_rows:
            # 剩余行数不足，返回当前c值
            return c_val
        # 取窗口内的c值（左闭右开，所以end_idx-1）
        window = df.loc[idx:end_idx-1, 'c']
    else:
        # 反向窗口：从当前索引开始往前window_size行
        start_idx = idx - window_size + 1
        if start_idx < 0:
            # 前面行数不足，返回当前c值
            return c_val
        # 取窗口内的c值
        window = df.loc[start_idx:idx, 'c']
    
    return window.mean()

# 应用函数到每一行
df['d'] = df.apply(lambda row: calculate_d(row, df), axis=1)

print(df)

运行后得到的结果和预期完全一致：

a  b    c      d
0  1  4  100  57.25
1  2 -2   53  76.50
2  3  2   48  38.00
3  4  0   28    NaN
4  5  1   18  18.00
5  6  4  202  202.00

备注：内容来源于stack exchange，提问作者narayanan_pr