为Pandas DataFrame按num_ID添加基于日期差阈值的total_times列

原始DataFrame #

需求 #

• 新增一列total_times
• 按num_ID分组，每个分组的首行total_times设为0
• 其余行判断：同分组上一行的end_date与当前行start_date间隔是否≥7天，是则设为1，否则设为0

中间状态示例 #

最终结果示例 #

解决方案代码 #

import pandas as pd

# 构造原始数据
data = {
    'num_ID': [1,1,2,2,2,3,3,4],
    'start_date': ['2022-02-14','2022-02-20','2022-03-11','2022-03-22','2022-04-10','2022-02-04','2022-02-14','2022-02-28'],
    'end_date': ['2022-02-15','2022-02-25','2022-03-21','2022-03-27','2022-04-15','2022-02-06','2022-02-23','2022-10-12'],
    'other_column': ['09:23:00','12:10:01','08:21:00','02:36:00','11:43:03','16:51:00','19:35:10','00:01:00']
}
df = pd.DataFrame(data)

# 转换日期列为datetime类型
df['start_date'] = pd.to_datetime(df['start_date'])
df['end_date'] = pd.to_datetime(df['end_date'])

# 分组处理，计算间隔并生成total_times
df['total_times'] = df.groupby('num_ID').apply(
    lambda x: pd.Series([0] + [(x['start_date'].iloc[i] - x['end_date'].iloc[i-1]).days >=7 for i in range(1, len(x))])
).reset_index(drop=True)

print(df)

代码解释 #

1. 日期类型转换：将start_date和end_date转为datetime类型，确保能正确计算日期差
2. 分组逻辑处理：按num_ID分组后，对每个分组单独操作：
   • 首行直接赋值0
   • 从第二行开始，计算当前行起始日期与上一行结束日期的天数差，判断是否≥7，返回1或0
3. 结果整合：将分组处理后的结果重置索引，赋值给新列total_times

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