Pandas按ID和Week两列分组 提取CloseContact前两个非空值生成新列

实现思路 #

• 按ID、Week两个字段对DataFrame执行分组操作
• 对每个分组内的CloseContact列做去空值处理，得到当前分组的有效接触记录序列
• 取有效序列的第一个值作为当前分组所有行的CC1值，取有效序列的第二个值作为当前分组所有行的CC2值，若对应位置无有效记录则填充空值
• 通过分组广播特性将两个字段的取值填充到当前分组的所有对应行中

实现代码 #

import pandas as pd
import numpy as np

NaN = np.nan
data = {'ID':['A','A','A','A','A','A','A','A','A','C','C','C','C','C','C','C','C'],
    'Week': ['Week1','Week1','Week1','Week1','Week2','Week2','Week2','Week2','Week3',
             'Week1','Week1','Week1','Week1','Week2','Week2','Week2','Week2'],
    'Risk':['High','','','','','','','','','High','','','','','','',''],
    'Testing':[NaN,'Pos',NaN,'Neg',NaN,NaN,NaN,NaN,'Pos', NaN, 
              NaN,NaN,'Negative',NaN,NaN,NaN,'Positive'],
    'CloseContact': [NaN, 'True', NaN, NaN, 'False',NaN, NaN, 'False', 'True', 
                    NaN, NaN, 'False', NaN, 'True','True','False', NaN ]}
    
df1 = pd.DataFrame(data)

# 按ID+Week分组
grouped = df1.groupby(['ID', 'Week'])['CloseContact']

# 生成CC1、CC2列
df1['CC1'] = grouped.transform(lambda x: x.dropna().iloc[0] if len(x.dropna()) >= 1 else pd.NA)
df1['CC2'] = grouped.transform(lambda x: x.dropna().iloc[1] if len(x.dropna()) >= 2 else pd.NA)

print(df1)

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