Pandas按分组条件从id_order列填充box列NaN值的方法

给定如下结构的Pandas DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'id_1': [19588.0, 16009.0, 19148.0, 16936.0, 13127.0, 15064.0, 10630.0, 12401.0, 13163.0, np.nan, np.nan, np.nan, 17577.0, 11808.0, 10538.0, 19527.0, 17697.0, 18426.0, 17632.0, 14726.0, 18361.0, 12908.0, 10279.0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    'box': [100.0, np.nan, np.nan, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0, np.nan, np.nan, 201.0, 202.0, 203.0, 204.0, 205.0, 206.0, 300.0, 301.0, 302.0, 303.0, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
    'reference': [300025886]*11 + [300025887]*9 + [300025888]*10,
    'id_2': [1.0, np.nan, np.nan, 3.0, 4.0, 6.0, 7.0, 8.0, 9.0, 5.0, 2.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 17.0, 18.0, 19.0, 20.0, np.nan, np.nan, 21.0, 22.0, 23.0, 24.0, np.nan, np.nan, np.nan],
    'id_order': [100.0, np.nan, np.nan, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0, 108.0, 107.0, 201.0, 202.0, 203.0, 204.0, 205.0, 206.0, 300.0, 301.0, 302.0, 303.0, np.nan, np.nan, 304.0, 305.0, 306.0, 307.0, np.nan, np.nan, np.nan]
})

需按reference分组，遵循以下规则生成new列：

• 规则1：同一reference下id_1非空数量与id_2非空数量相等时，用id_order中未被box使用的值填充box的NaN；
• 规则2：同一reference下id_1非空数量多于id_2非空数量时，不处理box的NaN；
• 规则3：同一reference下id_1非空数量少于id_2非空数量时，从id_order的空闲值中随机选取不重复值填充box的NaN；
• 所有规则下，new列优先保留box的非NaN值，仅填充符合条件的NaN。

核心思路是按reference分组后，针对每组统计关键指标，再匹配对应规则处理缺失值。

步骤1：定义分组处理函数 #

def process_group(group):
    # 统计当前组id_1、id_2的非空数量
    count_id1 = group['id_1'].notna().sum()
    count_id2 = group['id_2'].notna().sum()
    
    # 提取box已使用的值、id_order所有非空值，计算空闲候选值
    used_box = group['box'].dropna().unique()
    all_id_order = group['id_order'].dropna().unique()
    free_id_order = [x for x in all_id_order if x not in used_box]
    
    # 初始化new列
    group['new'] = group['box'].copy()
    # 获取需要填充的行索引
    fill_indices = group['new'].isna()
    fill_count = fill_indices.sum()

    if count_id1 == count_id2:
        # 规则1：按顺序用空闲值填充（确保候选值数量足够）
        if len(free_id_order) >= fill_count:
            group.loc[fill_indices, 'new'] = free_id_order[:fill_count]
    elif count_id1 < count_id2:
        # 规则3：随机选取不重复的空闲值填充
        if len(free_id_order) >= fill_count:
            np.random.shuffle(free_id_order)
            group.loc[fill_indices, 'new'] = free_id_order[:fill_count]
    # 规则2：id_1数量更多，不做处理
    
    return group

步骤2：分组应用函数 #

# 设置随机种子保证结果可复现（可选）
np.random.seed(42)

# 按reference分组执行处理逻辑
result_df = df.groupby('reference', group_keys=False).apply(process_group)

逻辑说明 #

1. 分组统计：每组内先计算id_1和id_2的非空数量，确定适用规则；
2. 候选池生成：从id_order中筛选出未被box使用的空闲值，作为填充候选；
3. 规则匹配：
   • 规则1：直接按顺序取空闲值填充缺失位置；
   • 规则3：先随机打乱候选列表，再取对应数量的值填充，保证随机性且无重复；
   • 规则2：跳过填充操作，保留原NaN；
4. 结果整合：处理后的各组拼接为最终DataFrame。

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