基于多条件IF逻辑填充Pandas DataFrame中的NaN值

输入数据 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "Set": [100, 100, 110, 110, 130, 130, 130, 140, 140, 150, 150, 150, 160, 170, 170],
    "measure": [np.nan, np.nan, 11, 10, np.nan, np.nan, np.nan, np.nan, np.nan, 10, 13, 8, np.nan, 12, 13],
    "width": [0.19, 0.18, 0.2, 0.27, 0.18, 0.17, 0.21, 0.19, 0.16, 0.19, 0.24, 0.3, 0.15, 0.32, 0.28]
})

填充规则（优先级从高到低） #

• 规则1：找到当前NaN所在Set与最近的非NaN邻居Set（绝对差最小），用该邻居Set的measure最大值填充NaN，同时Adjusted Width取该邻居对应measure最大值的width。
• 规则2：若存在两个邻居Set与当前Set的绝对差相同，则选择measure最大值更高的邻居Set，用其measure最大值填充，Adjusted Width取对应值。
• 规则3：若前两个条件均平局，则选择width最大值更高的邻居Set，用其对应measure值填充，Adjusted Width取该width值。

期望输出 #

Set  measure  width  Adjusted Width
0   100     11.0   0.19             0.2
1   100     11.0   0.18             0.2
2   110     11.0   0.20             0.2
3   110     10.0   0.27             0.27
4   130     13.0   0.18             0.24
5   130     13.0   0.17             0.24
6   130     13.0   0.21             0.24
7   140     13.0   0.19             0.24
8   140     13.0   0.16             0.24
9   150     10.0   0.19             0.19
10  150     13.0   0.24             0.24
11  150      8.0   0.30             0.30
12  160     13.0   0.15             0.28
13  170     12.0   0.32             0.32
14  170     13.0   0.28             0.28

实现代码 #

# 1. 预处理：提取有非NaN measure的Set的统计信息
valid_sets = df[df['measure'].notna()].groupby('Set').apply(
    lambda x: pd.Series({
        'max_measure': x['measure'].max(),
        'max_measure_width': x.loc[x['measure'].idxmax(), 'width']
    })
).reset_index()

# 获取所有有效Set的列表
valid_set_list = valid_sets['Set'].tolist()

# 2. 定义函数：根据当前Set选择最优邻居并返回填充值和Adjusted Width
def get_fill_values(current_set):
    # 计算当前Set与所有有效Set的绝对差
    valid_sets['abs_diff'] = abs(valid_sets['Set'] - current_set)
    
    # 按优先级排序：先按绝对差升序，再按max_measure降序，最后按max_measure_width降序
    sorted_neighbors = valid_sets.sort_values(
        by=['abs_diff', 'max_measure', 'max_measure_width'],
        ascending=[True, False, False]
    )
    
    # 取排序后的第一个邻居
    best_neighbor = sorted_neighbors.iloc[0]
    return best_neighbor['max_measure'], best_neighbor['max_measure_width']

# 3. 对需要填充的行应用函数
fill_mask = df['measure'].isna()
df.loc[fill_mask, ['measure', 'Adjusted Width']] = df.loc[fill_mask, 'Set'].apply(
    lambda x: pd.Series(get_fill_values(x))
).values

# 4. 对非NaN行，Adjusted Width取自身width
df.loc[~fill_mask, 'Adjusted Width'] = df.loc[~fill_mask, 'width']

# 调整列顺序（与期望输出一致）
df = df[['Set', 'measure', 'width', 'Adjusted Width']]

print(df.to_string(index=False))

代码说明 #

1. 预处理步骤：通过groupby计算每个有非NaN值的Set的measure最大值，以及该最大值对应的width，得到有效邻居的统计信息。
2. 邻居选择函数：对每个需要填充的Set，计算其与所有有效Set的绝对差，然后按照规则的优先级排序，选出最优的邻居。
3. 填充与赋值：对所有NaN行应用选择函数，填充measure和Adjusted Width；非NaN行直接使用自身的width作为Adjusted Width。
4. 格式调整：最后调整列顺序，确保输出与期望一致。

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