如何在Python 3.8中为DataFrame创建变量存储同一id下被照护者的年龄

我来帮你搞定Python 3.8环境下的这个需求，先把数据和需求理清楚：

数据初始化 #

首先是你提供的DataFrame数据：

import pandas as pd
import numpy as np

df = pd.DataFrame({'id': [1,1,2,2,2,3,3,4,4,4,5,5,6,6], 
                   'rank' : [1,2,1,2,3,1,2,1,2,3,1,2,1,2], 
                   'age' : [20,2,30,28,4,26,3,22,23,1,34,29,56,25], 
                   'sex' : ['female','female','male','female','male','female','male','female','male','male','male','female','female','female'], 
                   'care' : [np.nan,np.nan,3,np.nan,np.nan,2,np.nan,np.nan,3,np.nan,np.nan,np.nan,2,np.nan]})

变量说明 #

数据包含以下字段：

• id：家庭识别编号
• rank：家庭成员的排名（1=户主，2=配偶，3=子女）
• age：人员年龄
• sex：人员性别
• care：接受照护服务的家庭成员的rank值

核心需求 #

创建名为new的变量，存储同一家庭（同一id）内接受照护服务人员的年龄。你已经在Stata和R中实现了这个逻辑，下面提供三种Python 3.8兼容的实现方案：

方法一：分组映射法（贴近R代码逻辑） #

这个方法和你在R中用group_by + mutate的思路完全一致，按家庭分组后，为每个组构建rank到age的映射，再用care匹配对应年龄：

def get_care_age(group):
    # 为当前家庭创建rank到age的字典映射
    rank_age_map = group.set_index('rank')['age'].to_dict()
    # 用care字段匹配对应的年龄，无匹配则保留NaN
    group['new'] = group['care'].map(rank_age_map)
    return group

# 分组处理后重置索引
df = df.groupby('id').apply(get_care_age).reset_index(drop=True)

方法二：全局字典匹配法 #

如果数据量不大，可以先构建一个全局的(id, rank)到age的映射字典，再逐行匹配：

# 构建(id, rank) -> age的全局映射
id_rank_age = df.set_index(['id', 'rank'])['age'].to_dict()

# 生成new变量，找不到匹配时返回NaN
df['new'] = df.apply(lambda row: id_rank_age.get((row['id'], row['care']), np.nan), axis=1)

方法三：合并关联法 #

通过创建临时表提取照护对象的年龄，再合并回原数据，适合熟悉SQL关联逻辑的场景：

# 先筛选有care值的行，将rank重命名为care，提取对应年龄作为new
care_age_df = df.dropna(subset=['care']).rename(columns={'rank': 'care', 'age': 'new'})[['id', 'care', 'new']]

# 左连接合并回原数据，保留所有行
df = df.merge(care_age_df, on=['id', 'care'], how='left')

最终输出效果 #

运行任意一种方法后，你都会得到包含new变量的DataFrame，打印结果如下：

id  rank  age     sex  care   new
0    1     1   20  female   NaN   NaN
1    1     2    2  female   NaN   NaN
2    2     1   30    male   3.0   4.0
3    2     2   28  female   NaN   NaN
4    2     3    4    male   NaN   NaN
5    3     1   26  female   2.0   3.0
6    3     2    3    male   NaN   NaN
7    4     1   22  female   NaN   NaN
8    4     2   23    male   3.0   1.0
9    4     3    1    male   NaN   NaN
10   5     1   34    male   NaN   NaN
11   5     2   29  female   NaN   NaN
12   6     1   56  female   2.0  25.0
13   6     2   25  female   NaN   NaN

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