如何基于字典高效填充DataFrame的指定列

问题背景 #

现有如下DataFrame和字典：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'A': [1, 1, 1, 2, 2, 3, 3, 3, 3],
    'ignore_me': range(9),
    'fill_me': [np.nan] * 9
})

di = {
    1: ['a', 'b'],
    2: ['c', 'd'],
    3: ['e', 'f', 'g']
}

需求是：根据A列的值匹配字典di的键，用对应字典值列表填充fill_me列，每个组仅填充前N行（N为对应字典值的长度），超出部分保留NaN，预期结果如下：

A  ignore_me fill_me
0  1          0       a
1  1          1       b
2  1          2     NaN
3  2          3       c
4  2          4       d
5  3          5       e
6  3          6       f
7  3          7       g
8  3          8     NaN

原实现方式需要循环、concat且生成多个临时DataFrame：

df_list = []
for key, values in di.items():
    temp_df = df[df['A'] == key].reset_index(drop=True)
    fill_df = pd.DataFrame({'A': [key]* len(values), 'fill_me': values})
    df_list.append(temp_df.combine_first(fill_df))

final_df = pd.concat(df_list, ignore_index=True)

希望找到更简洁的实现方式，最好能直接在原DataFrame上操作。

简洁实现方案 #

方案一：组内序号匹配赋值（直接修改原df） #

无需额外生成DataFrame，直接在原数据上操作：

# 生成每个组内的连续序号
df['group_idx'] = df.groupby('A').cumcount()

# 把字典转换成(A, 组内序号)对应值的映射表
fill_map = {
    (key, idx): val 
    for key, vals in di.items() 
    for idx, val in enumerate(vals)
}

# 匹配赋值，没匹配到的保留原NaN
df['fill_me'] = df.apply(
    lambda row: fill_map.get((row['A'], row['group_idx']), row['fill_me']),
    axis=1
)

# 清理临时列
df = df.drop('group_idx', axis=1)

方案二：构造匹配DataFrame合并填充 #

通过生成一个包含匹配规则的小DataFrame，和原数据合并完成填充：

# 把字典展开成带组内序号的填充DataFrame
fill_df = pd.DataFrame([
    (key, idx, val) 
    for key, vals in di.items() 
    for idx, val in enumerate(vals)
], columns=['A', 'group_idx', 'fill_me'])

# 给原df添加组内序号
df['group_idx'] = df.groupby('A').cumcount()

# 左合并保留所有行，用新值填充原NaN
df = df.merge(fill_df, on=['A', 'group_idx'], how='left', suffixes=('', '_new'))
df['fill_me'] = df['fill_me_new'].combine_first(df['fill_me'])

# 清理临时列
df = df.drop(['group_idx', 'fill_me_new'], axis=1)

方案三：向量化高效操作（推荐） #

避免apply的逐行循环，用布尔索引和映射实现更高效的填充：

# 生成组内序号
df['group_idx'] = df.groupby('A').cumcount()

# 为每个A获取可填充的最大序号（字典值长度-1）
max_valid_idx = df['A'].map(lambda x: len(di[x]) - 1)

# 筛选出需要填充的行
fill_mask = df['group_idx'] <= max_valid_idx

# 对需要填充的行，直接从字典中取值
df.loc[fill_mask, 'fill_me'] = df.loc[fill_mask].apply(
    lambda x: di[x['A']][x['group_idx']],
    axis=1
)

# 清理临时列
df = df.drop('group_idx', axis=1)

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