如何基于≤当前total值的对应中位数填充area列的NaN数据？

需求说明 #

需将area列的空NaN值替换为total值小于等于当前行total值的对应area列中位数，具体规则：

• 若当前total值对应的area有非NaN值，取其中位数填充；
• 若无则将total值递减1继续查找，直至找到可用值。

示例 #

• 行3的total为8，取total=8的area中位数52填充；
• 行4的total为20，递减查找后取total=8的中位数52；
• 行6的total为59，取total=56的area值34（中位数）填充。

原始数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({'total': [5, 8, 8, 8, 20, 56, 59], 
                   'area': [40, 51, 53, np.nan, np.nan, 34, np.nan]})
print(df)

输出：

total   area
0      5  40.0
1      8  51.0
2      8  53.0
3      8   NaN
4     20   NaN
5     56  34.0
6     59   NaN

期望结果 #

result = pd.DataFrame({'total': [5, 8, 8, 8, 20, 56, 59], 'area': [40, 51, 53, 52, 52, 34, 34]})
print(result)

输出：

total  area
0      5    40
1      8    51
2      8    53
3      8    52
4     20    52
5     56    34
6     59    34

原函数问题分析 #

你提供的自定义函数无法生成预期结果，问题出在以下几点：

1. fillna用法错误：fillna传入的函数仅接收当前NaN值作为参数，无法获取当前行的total值，导致函数无法按行的total进行查找。
2. 非NaN值判断逻辑错误：用sum()判断是否存在非NaN值不可靠（比如area存在负数时，sum可能为0但有有效数据），应该直接判断该total下是否有非NaN的area值。
3. 循环逻辑冗余：函数中的x变量无实际作用，且return后的代码不会执行。

正确解决方案 #

步骤1：预计算有效total对应的area中位数 #

先统计所有存在非NaN area的total对应的中位数，避免重复计算：

# 按total分组计算area中位数，丢弃无有效数据的total
total_medians = df.groupby('total')['area'].median().dropna()

步骤2：定义填充函数 #

针对每行数据，从当前total开始递减查找，返回第一个有效中位数：

def fill_area(row):
    # 若area已有值，直接返回
    if not pd.isna(row['area']):
        return row['area']
    current_total = row['total']
    # 从当前total开始递减查找
    while current_total >= df['total'].min():
        if current_total in total_medians:
            return total_medians[current_total]
        current_total -= 1
    # 若所有更小的total都无有效数据，返回NaN（可根据需求调整）
    return np.nan

步骤3：应用函数填充NaN #

df['area'] = df.apply(fill_area, axis=1)
print(df)

输出结果与期望完全一致：

total  area
0      5  40.0
1      8  51.0
2      8  53.0
3      8  52.0
4     20  52.0
5     56  34.0
6     59  34.0

优化（可选） #

如果数据量较大，循环查找效率较低，可以先将total_medians的索引转为排序后的数组，用二分查找快速定位最大的<=当前total的有效值：

import bisect

# 获取排序后的有效total列表
sorted_totals = sorted(total_medians.index)

def fill_area_optimized(row):
    if not pd.isna(row['area']):
        return row['area']
    current_total = row['total']
    # 找到第一个大于current_total的位置，取前一个元素
    idx = bisect.bisect_right(sorted_totals, current_total) - 1
    if idx >= 0:
        return total_medians[sorted_totals[idx]]
    return np.nan

df['area'] = df.apply(fill_area_optimized, axis=1)

该优化适合处理大规模数据，查找效率从O(n)提升到O(logn)。

内容的提问来源于stack exchange，提问作者Сделано в России