如何基于含NaN的Pandas DataFrame df1构造目标DataFrame df2？

问题背景 #

现有如下构造的Pandas DataFrame df1：

import pandas as pd
import numpy as np

data={'ID':[1,1,1,1,1,2,2,2],
'city':['A',0,0,'C','D','E',0,'F']}
df1=pd.DataFrame(data)
df1.replace(0,np.nan,inplace=True)

df1按ID分组后，city列存在NaN值。期望得到目标DataFrame df2，构造代码如下：

import pandas as pd
import numpy as np

data={'ID':[1,1,1,1,1,1,1,2,2,2],
'city':['A','C','D','C','D','C','D','E','F','F']}
df2=pd.DataFrame(data)
df2.replace(0,np.nan,inplace=True)

需求为：每个ID分组内，用组内非NaN的city值填充原NaN行并新增对应行。尝试bfill()方法仅填充NaN但未新增行，不符合需求，需实现从df1到df2的转换。

实现方案 #

核心逻辑是先提取每个分组的非NaN城市列表，再针对原DataFrame的每一行处理：非NaN行直接保留，NaN行则用分组内所有非NaN城市生成新行。具体代码如下：

步骤1：预存每个分组的非NaN城市集合 #

# 按ID分组，提取每组的非NaN城市去重列表
city_groups = df1.groupby('ID')['city'].apply(lambda x: x.dropna().unique().tolist()).to_dict()

步骤2：生成目标行数据 #

new_data = []
for _, row in df1.iterrows():
    if pd.notna(row['city']):
        # 非NaN行直接加入结果
        new_data.append(row.to_dict())
    else:
        # NaN行，遍历对应分组的所有非NaN城市，生成新行
        for city in city_groups[row['ID']]:
            new_data.append({'ID': row['ID'], 'city': city})

步骤3：转换为目标DataFrame #

df2 = pd.DataFrame(new_data)

执行上述代码后，得到的df2与目标构造的结果完全一致。

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