Python按ID分组获取首行值问题：需保留NaN的解决方案

当前代码通过分组聚合获取team列的首个值时，会自动忽略NaN，取分组内第一个非空值，但实际需求是保留分组内第一行的原始team值（即使是NaN），并将其映射到对应ID的所有行的first_team列。

样本数据集（当前错误输出） #

ID     date           name       team       first_team
 101   05/2012         James      NaN            NY
 101   07/2012         James      NY             NY
 102   06/2013         Adams      NC             NC
 102   05/2014         Adams      AL             NC 

现有代码 #

first_dict = df.groupby('ID').agg({'team':'first'}).to_dict()['team']
df['first_team'] = df['ID'].apply(lambda x: first_dict[x])

期望输出 #

ID      date        name      team         first_team 
  101     05/2012     James      NaN           NaN 
  101     07/2012     James      NY            NaN 
  102     06/2013     Adams      NC            NC 
  102     05/2014     Adams      AL            NC 

解决方案 #

问题出在agg('first')会自动跳过NaN，我们需要用nth(0)来获取分组内的第一行原始值（包含NaN），修改后的代码如下：

# 获取每个ID分组的第一行team值（保留NaN）
first_dict = df.groupby('ID')['team'].nth(0).to_dict()
df['first_team'] = df['ID'].map(first_dict)

说明 #

• groupby('ID')['team'].nth(0)会严格返回每个分组第0个位置的原始值，不会自动过滤NaN；
• 用map()替代apply(lambda x: ...)，语法更简洁，性能也更优。

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