计算索引日期与首次终止标识日期的天数差（无标识取末次观测）

问题 #

现有包含id、随访日期date、分组group、终止标识indicator（值为1代表终止）的DataFrame，需实现以下需求：

• 对对应行填入索引日期（每个id+group组内的最早日期）与首次出现终止标识日期的天数差；
• 若组内未出现终止标识，则填入索引日期至组内末次观测日期的天数；
• 仅在终止发生的行（indicator=1）或无终止时的末次行填充该值，其余行填NA。

原始数据 #

id  date             group   indicator      
1   15-01-2022        1       0
1   15-01-2022        2       0
1   16-01-2022        2       1
1   20-01-2022        2       0
2   18-01-2022        1       0
2   20-01-2022        2       0
2   27-01-2022        2       0

期望结果 #

id  date             group   indicator     stoptime
1   15-01-2022        1       0             NA
1   15-01-2022        2       0             NA
1   16-01-2022        2       1             1
1   20-01-2022        2       0             NA
2   18-01-2022        1       0             NA
2   20-01-2022        2       0             NA
2   27-01-2022        2       0             9

解决方案 #

以下是基于Pandas的实现方案，核心逻辑是按id+group分组计算关键日期，再匹配填充目标列：

步骤1：数据预处理，转换日期格式 #

先将date列转为datetime类型，方便后续日期计算：

import pandas as pd

# 构造原始数据
df = pd.DataFrame({
    'id': [1,1,1,1,2,2,2],
    'date': ['15-01-2022','15-01-2022','16-01-2022','20-01-2022','18-01-2022','20-01-2022','27-01-2022'],
    'group': [1,2,2,2,1,2,2],
    'indicator': [0,0,1,0,0,0,0]
})

# 转换日期格式（适配日-月-年的输入格式）
df['date'] = pd.to_datetime(df['date'], format='%d-%m-%Y')

步骤2：按分组计算关键日期 #

按id和group分组，提取每组的三个核心日期：

# 分组计算索引日期、首次终止日期、末次观测日期
grouped_info = df.groupby(['id', 'group']).agg(
    index_date=('date', 'min'),
    first_stop_date=('date', lambda x: x[df.loc[x.index, 'indicator'] == 1].min()),
    last_date=('date', 'max')
).reset_index()

步骤3：合并数据并计算stoptime #

将分组信息合并回原数据，根据规则填充stoptime：

# 合并分组信息到原DataFrame
df = df.merge(grouped_info, on=['id', 'group'], how='left')

# 初始化stoptime列为NA
df['stoptime'] = pd.NA

# 填充终止行的天数差：首次终止日期 - 索引日期
stop_row_mask = df['indicator'] == 1
df.loc[stop_row_mask, 'stoptime'] = (df.loc[stop_row_mask, 'first_stop_date'] - df.loc[stop_row_mask, 'index_date']).dt.days

# 标记每个分组的末次行
df['is_group_last'] = df.groupby(['id', 'group'])['date'].rank(ascending=False, method='first') == 1

# 填充无终止组的末次行天数差：末次观测日期 - 索引日期
no_stop_last_mask = df['first_stop_date'].isna() & df['is_group_last']
df.loc[no_stop_last_mask, 'stoptime'] = (df.loc[no_stop_last_mask, 'last_date'] - df.loc[no_stop_last_mask, 'index_date']).dt.days

# 清理临时辅助列
df = df.drop(columns=['index_date', 'first_stop_date', 'last_date', 'is_group_last'])

# 输出结果
print(df.to_string(index=False))

运行结果 #

输出与期望结果一致：

id       date  group  indicator  stoptime
 1 2022-01-15      1          0       <NA>
 1 2022-01-15      2          0       <NA>
 1 2022-01-16      2          1         1
 1 2022-01-20      2          0       <NA>
 2 2022-01-18      1          0       <NA>
 2 2022-01-20      2          0       <NA>
 2 2022-01-27      2          0         9

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