如何在DataFrame中按Subject分组后添加对应ScreeningDate列

给定如下DataFrame数据：

{'index': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 49, 50, 51, 52, 53, 54, 55, 56, 57],
 'columns': ['Subject', 'Visit', 'Date'],
 'data': [['A', 'Screening', pd.Timestamp('2023-11-15 00:00:00')],
  ['A', 'Week 0', pd.Timestamp('2023-11-29 00:00:00')],
  ['A', 'Week 2', pd.Timestamp('2023-12-12 00:00:00')],
  ['A', 'Week 4', pd.Timestamp('2023-12-27 00:00:00')],
  ['A', 'Week 8', pd.Timestamp('2024-01-22 00:00:00')],
  ['A', 'Week 12', pd.Timestamp('2024-02-21 00:00:00')],
  ['A', 'Week 16', pd.Timestamp('2024-03-17 00:00:00')],
  ['A', 'Week 20', pd.Timestamp('2024-04-17 00:00:00')],
  ['A', 'Week 28', pd.Timestamp('2024-06-06 00:00:00')],
  ['A', 'Week 36', pd.Timestamp('2024-08-08 00:00:00')],
  ['B', 'Screening', pd.Timestamp('2024-02-19 00:00:00')],
  ['B', 'Week 0', pd.Timestamp('2024-03-10 00:00:00')],
  ['B', 'Week 2', pd.Timestamp('2024-03-24 00:00:00')],
  ['B', 'Week 4', pd.Timestamp('2024-04-07 00:00:00')],
  ['B', 'Week 8', pd.Timestamp('2024-05-05 00:00:00')],
  ['B', 'Week 12', pd.Timestamp('2024-06-02 00:00:00')],
  ['B', 'Week 16', pd.Timestamp('2024-06-27 00:00:00')],
  ['B', 'Week 20', pd.Timestamp('2024-07-28 00:00:00')],
  ['B', 'Week 28', pd.Timestamp('2024-09-04 00:00:00')]],
  'index_names': [None],
  'column_names': [None]}

需要为该DataFrame新增一列ScreeningDate，填充每个Subject对应的Screening日期。

方法一：使用groupby + transform #

先将原始数据转换为DataFrame（如果尚未转换），然后通过分组筛选的方式填充新列：

import pandas as pd

# 转换为DataFrame
df = pd.DataFrame(
    data=your_data['data'],
    columns=your_data['columns'],
    index=your_data['index']
)

# 新增ScreeningDate列
df['ScreeningDate'] = df.groupby('Subject')['Date'].transform(
    lambda x: x[df.loc[x.index, 'Visit'] == 'Screening'].iloc[0]
)

该方法按Subject分组后，对每组的Date列执行筛选，提取出对应Visit为Screening的日期，再将这个值填充到组内所有行。

方法二：构建映射字典（更高效） #

先创建Subject到其Screening日期的映射关系，再通过映射快速填充新列，适合数据量较大的场景：

import pandas as pd

# 转换为DataFrame
df = pd.DataFrame(
    data=your_data['data'],
    columns=your_data['columns'],
    index=your_data['index']
)

# 构建Subject与Screening日期的映射字典
screening_map = df[df['Visit'] == 'Screening'].set_index('Subject')['Date'].to_dict()

# 填充新列
df['ScreeningDate'] = df['Subject'].map(screening_map)

该方法先筛选出所有Screening记录，将Subject设为索引后转为字典，再用map方法将每个Subject对应的日期批量填充到新列中，执行效率更高。

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