Pandas Groupby按分组基于日期差计算多个自定义字段的实现方案

步骤1：数据预处理 #

首先导入依赖库，构造测试数据集，将日期列转换为标准datetime格式，保证后续日期计算的正确性。

步骤2：分组逻辑实现 #

按CID字段拆分独立分组，对每个分组依次完成基准重置标记、A列、B列、days_diff列的计算：

• 标记所有需要重置基准的位置：分组首行默认重置，上一行MID非空时当前行重置，当前行与上一基准日期差值超过30天时当前行重置
• A列由重置标记的累计和直接得到，初始值为1，每次重置计数+1
• B列按A列分组后取每组第一行的RefID值，同一A值的行B值保持一致
• days_diff列按A列分组，计算当前日期与分组首行基准日期的差值，每组首行的差值设为NaT

完整代码 #

import pandas as pd
import numpy as np

# 构造测试数据
data = [
    [100, 1, '1/01/2021', 'A', np.nan],
    [100, 2, '3/01/2021', 'A', np.nan],
    [100, 3, '4/01/2021', 'A', 101],
    [100, 4, '15/01/2021', 'A', np.nan],
    [100, 5, '18/01/2021', 'A', np.nan],
    [200, 6, '3/03/2021', 'B', np.nan],
    [200, 7, '4/04/2021', 'B', np.nan],
    [200, 8, '9/04/2021', 'B', 102],
    [200, 9, '25/04/2021', 'B', np.nan],
    [300, 10, '26/04/2021', 'C', np.nan],
    [300, 11, '27/05/2021', 'C', np.nan],
    [300, 12, '28/05/2021', 'C', 103]
]
df = pd.DataFrame(data, columns=['CID', 'RefID', 'Date', 'Group', 'MID'])
# 转换日期格式为日/月/年
df['Date'] = pd.to_datetime(df['Date'], format='%d/%m/%Y')

# 单个CID分组的处理函数
def process_cid_group(g):
    g = g.reset_index(drop=True)
    # 初始化重置标记：首行默认重置基准
    reset_flags = [True]
    current_base_date = g.loc[0, 'Date']
    # 遍历剩余行判断是否需要重置
    for i in range(1, len(g)):
        # 重置条件1：上一行MID不为空
        prev_mid_valid = pd.notna(g.loc[i-1, 'MID'])
        # 重置条件2：当前日期与基准日期差值超过30天
        date_diff_over30 = (g.loc[i, 'Date'] - current_base_date).days > 30
        if prev_mid_valid or date_diff_over30:
            reset_flags.append(True)
            current_base_date = g.loc[i, 'Date']
        else:
            reset_flags.append(False)
    # 计算A列
    g['A'] = np.cumsum(reset_flags)
    # 计算B列
    g['B'] = g.groupby('A')['RefID'].transform('first')
    # 计算days_diff列
    g['days_diff'] = g['Date'] - g.groupby('A')['Date'].transform('first')
    # 每组首行差值设为NaT
    g.loc[g.groupby('A').head(1).index, 'days_diff'] = pd.NaT
    return g

# 按CID分组应用处理逻辑
df = df.groupby('CID', group_keys=False).apply(process_cid_group)
# 调整列顺序匹配预期输出
df = df[['CID', 'RefID', 'Date', 'Group', 'MID', 'days_diff', 'A', 'B']]
# 输出结果
print(df)

运行代码得到的结果与你提供的预期输出完全一致。

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