如何按后续非NaN值均分规则填充DataFrame的NaN缺失值？

问题背景 #

现有如下DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'columnA': ['A', 'A', 'A', 'B', 'B', 'A', 'D'],
    'columnB': ['A', 'B', 'C', 'A', 'C', 'D', 'C'],
    'columnC': [10, np.nan, 20, 30, np.nan, np.nan, 15]
})

需要按规则填充columnC的NaN值：找到每个NaN区间后续的首个非NaN值，用该值除以区间内的总条目数（包括这个非NaN值本身），结果填充整个区间的所有行。

期望输出：

# 填充后的结果
  columnA columnB  columnC
0       A       A     10.0
1       A       B     10.0
2       A       C     10.0
3       B       A     30.0
4       B       C      5.0
5       A       D      5.0
6       D       C      5.0

补充说明：示例里20除以2（对应2行）得10，填充第1、2行；15除以3（对应3行）得5，填充第4、5、6行。

解决代码 #

# 1. 给每个需要统一填充的区间标记分组
df['group'] = df['columnC'].notna()[::-1].cumsum()[::-1]

# 2. 计算每个分组的填充值：组内最后一个非NaN值 ÷ 组行数
fill_values = df.groupby('group')['columnC'].apply(lambda x: x.dropna().iloc[-1] / len(x))

# 3. 映射填充值到原列
df['columnC'] = df['group'].map(fill_values)

# 移除辅助列
df = df.drop('group', axis=1)

print(df)

步骤说明 #

• 分组标记：把每个非NaN值和它前面连续的NaN值划成一组。反向累计求和的方式能准确捕捉每个“NaN区间+后续首个非NaN值”的范围。
• 计算填充值：对每个分组，取组内唯一的非NaN值（也就是区间末尾的那个值），除以组内的总行数，得到该组所有行的填充值。
• 映射填充：把每个分组对应的填充值批量替换到columnC列，完成填充。

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