基于反向移动平均填充DataFrame中NaN值的高效实现方案

问题描述 #

现有如下带NaN值的DataFrame：

A   B   C
0   NaN NaN NaN
1   NaN NaN NaN
2   2.0 1.0 NaN
3   2.0 NaN NaN
4   0.0 4.0 2.0
5   NaN 2.0 5.0
6   NaN 3.0 1.0

需要填充C列所有NaN值，填充规则为反向递归取后续3个值的移动平均：

• 行3的C值 = (行4的2.0 + 行5的5.0 + 行6的1.0)/3 = 2.666667
• 行2的C值 = (行3填充后的2.666667 + 行4的2.0 + 行5的5.0)/3 = 3.222222
• 需完全填充无NaN，且全量数据集达3万行（含20%缺失值），需高效实现。

核心问题 #

普通pd.rolling()结合反转/位移的方法无效，因为这类方法仅基于原始数据的窗口计算，无法处理递归填充——即前面的NaN填充值依赖后面已经填充好的结果，而非原始的NaN。纯Python循环效率过低，无法应对3万行规模。

高效解决方案 #

使用numba编译循环实现反向递归计算，速度比纯Python循环快数十倍，完全适配大数据量：

步骤1：构造示例数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'A': [np.nan, np.nan, 2.0, 2.0, 0.0, np.nan, np.nan],
    'B': [np.nan, np.nan, 1.0, np.nan, 4.0, 2.0, 3.0],
    'C': [np.nan, np.nan, np.nan, np.nan, 2.0, 5.0, 1.0]
})

步骤2：用numba实现递归填充 #

from numba import jit

@jit(nopython=True)
def fill_backward_rolling_mean(arr, window=3):
    n = len(arr)
    # 从后往前遍历，优先处理有足够后续值的行
    for i in range(n - window, -1, -1):
        if np.isnan(arr[i]):
            end_idx = min(i + window, n)
            arr[i] = np.nanmean(arr[i+1:end_idx])
    # 处理最前面不足window的行
    for i in range(n - window - 1, -1, -1):
        if np.isnan(arr[i]):
            end_idx = min(i + window, n)
            arr[i] = np.nanmean(arr[i+1:end_idx])
    return arr

# 复制C列避免修改原始数据
c_values = df['C'].copy().values
# 执行填充
filled_c = fill_backward_rolling_mean(c_values, window=3)
# 赋值回DataFrame
df['C'] = filled_c

步骤3：验证结果 #

填充后的C列与目标完全一致：

C
0  2.839506
1  2.629630
2  3.222222
3  2.666667
4  2.000000
5  5.000000
6  1.000000

备选方案（无numba依赖） #

如果无法安装numba，可使用纯pandas反向遍历，速度略慢但3万行仍可接受：

c_col = df['C'].copy()
window = 3
# 从后往前遍历填充
for i in range(len(c_col)-2, -1, -1):
    if pd.isna(c_col.iloc[i]):
        window_vals = c_col.iloc[i+1:i+1+window]
        c_col.iloc[i] = window_vals.mean()
df['C'] = c_col

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