Pandas中按ID分组结合状态与日期阈值生成标记列的问题

需求说明 #

按ID分组后，执行以下逻辑生成flag列：

• 找到每组中**第一个status为'S'**的日期作为初始基准日期
• 后续行与当前基准日期比较：
  • 若日期差≥30天且当前行status为'S'：将该行日期设为新基准，flag标记为当前组号+1
  • 若日期差<30天：无论status如何，flag标记为当前组号
  • 不满足上述条件的行（如日期差≥30天但status不为'S'、组内无任何'S'行）：flag标记为0

初始测试数据 #

import pandas as pd
data = {"ID": [117, 117, 117, 117, 117, 117, 118, 118, 118, 118, 118, 118], 
        "Date": ["2023-11-14", "2024-01-25", "2024-02-01", "2024-02-04", "2024-02-11", "2024-03-04",
        "2024-01-02", "2024-01-28", "2024-02-04", "2024-02-18", "2024-03-11", "2024-06-05"], 
        "status": ['S', 'S', 'S', 'E', 'E', 'E', 'E', 'E', 'S', 'S', 'S', 'E']}
df = pd.DataFrame(data)

预期输出示例 #

• 患者117：index0（status为S）标记1；index1（超30天且status为S）标记2；index2-4（30天内）标记2；index5（超30天但非S）标记0。
• 患者118：index6-7（无S）标记0；index8（首个S）标记1；index9（30天内）标记1；index10（超30天且S）标记2；index11（超30天非S）标记0。

当前代码的问题 #

现有代码仅按日期阈值分组标记，未正确判断新基准行的status是否为'S'，无法满足需求：

import pandas as pd
data = {"ID": [117, 117, 117, 117, 117, 117, 118, 118, 118, 118, 118, 118], 
        "Date": ["2023-11-14", "2024-01-25", "2024-02-01", "2024-02-04", "2024-02-11", "2024-03-04",
        "2024-01-02", "2024-01-28", "2024-02-04", "2024-02-18", "2024-03-11", "2024-06-05"], 
        "status": ['S', 'S', 'S', 'E', 'E', 'E', 'E', 'E', 'S', 'S', 'S', 'E']}
df = pd.DataFrame(data)

# 自定义函数
def get_flag(d, thresh=30):
    dates = pd.to_datetime(d['Date'])
    status = d['status']
    ref = dates.iloc[0]
    result = [1]
    n = 2
    e = 2
    
    for date in dates.iloc[1:]:
        if (date - ref).days >= thresh:
            result.append(n)
            ref = date 
            n+=1
        else:
            result.append(e)
    return d.assign(flag=result)
        
# 分组应用
out = df.groupby('ID', group_keys=False).apply(get_flag)
out

更新后的测试数据与问题 #

修改后的测试数据如下：

import pandas as pd
data = {"ID": [117, 117, 117, 117, 117, 117, 117, 117,117,117,117,118, 118, 118, 118, 118, 118], 
        "Date": ["2013-09-18", "2016-02-07", "2016-02-17", "2016-02-20", "2016-04-05", "2016-04-12", "2016-04-12", "2016-04-14",  "2016-04-16", "2016-05-05", "2016-05-16",
        "2024-01-02", "2024-01-28", "2024-02-04", "2024-02-18", "2024-03-11", "2024-03-12"], 
        "status": ['E', 'E', 'B', 'E', 'E', 'S', 'B', 'S', 'E', 'E', 'E', 
                   'E', 'S', 'E', 'S', 'E', 'S']}
df = pd.DataFrame(data)

预期输出 #

• 患者117：index5（status为S）标记1，index6-9（30天内）标记1；其余行标记0
• 患者118：index12（status为S）标记1，index13-14（30天内）标记1，index16（超30天且S）标记2；其余行标记0

但使用以下函数后所有行标记为0，不符合预期：

def get_flag(g, thresh=30):
    dates = pd.to_datetime(g['Date'])
    status = g['status'].eq('S').astype(int)
    ref = dates.iloc[0]
    s = status.iloc[0]
    result = [s]

    for i in range(1, len(g)):
        date = dates.iloc[i]
        stat = status.iloc[i]
        if (date - ref).days >= thresh:
            result.append(result[-1]+stat if stat else 0)
            ref = date 
        else:
            result.append(result[-1])
    return g.assign(flag=result)
        
# 分组应用
out = df.groupby('ID', group_keys=False).apply(get_flag)

修正后的代码 #

import pandas as pd

def get_flag(group, thresh=30):
    # 转换日期格式
    dates = pd.to_datetime(group['Date'])
    status = group['status']
    # 初始化结果列表，默认全为0
    result = [0] * len(group)
    # 找到当前组第一个status为'S'的索引
    first_s_idx = status[status == 'S'].index.min()
    
    if pd.isna(first_s_idx):
        # 组内无'S'行，直接返回全0
        return group.assign(flag=result)
    
    # 初始基准日期和组号
    current_ref = dates.loc[first_s_idx]
    current_group = 1
    # 标记第一个'S'行
    result[group.index.get_loc(first_s_idx)] = current_group
    
    # 遍历第一个'S'之后的行
    for idx in group.index[group.index > first_s_idx]:
        date_diff = (dates.loc[idx] - current_ref).days
        current_status = status.loc[idx]
        
        if date_diff >= thresh:
            if current_status == 'S':
                # 更新基准和组号，标记当前行
                current_group += 1
                current_ref = dates.loc[idx]
                result[group.index.get_loc(idx)] = current_group
            else:
                # 日期超阈值但非'S'，标记0
                result[group.index.get_loc(idx)] = 0
        else:
            # 日期在阈值内，沿用当前组号
            result[group.index.get_loc(idx)] = current_group
    
    # 处理第一个'S'之前的行，全部标记0
    for idx in group.index[group.index < first_s_idx]:
        result[group.index.get_loc(idx)] = 0
    
    return group.assign(flag=result)

# 测试初始数据
data_initial = {"ID": [117, 117, 117, 117, 117, 117, 118, 118, 118, 118, 118, 118], 
        "Date": ["2023-11-14", "2024-01-25", "2024-02-01", "2024-02-04", "2024-02-11", "2024-03-04",
        "2024-01-02", "2024-01-28", "2024-02-04", "2024-02-18", "2024-03-11", "2024-06-05"], 
        "status": ['S', 'S', 'S', 'E', 'E', 'E', 'E', 'E', 'S', 'S', 'S', 'E']}
df_initial = pd.DataFrame(data_initial)
out_initial = df_initial.groupby('ID', group_keys=False).apply(get_flag)
print("初始数据输出：")
print(out_initial)

# 测试更新后的数据
data_updated = {"ID": [117, 117, 117, 117, 117, 117, 117, 117,117,117,117,118, 118, 118, 118, 118, 118], 
        "Date": ["2013-09-18", "2016-02-07", "2016-02-17", "2016-02-20", "2016-04-05", "2016-04-12", "2016-04-12", "2016-04-14",  "2016-04-16", "2016-05-05", "2016-05-16",
        "2024-01-02", "2024-01-28", "2024-02-04", "2024-02-18", "2024-03-11", "2024-03-12"], 
        "status": ['E', 'E', 'B', 'E', 'E', 'S', 'B', 'S', 'E', 'E', 'E', 
                   'E', 'S', 'E', 'S', 'E', 'S']}
df_updated = pd.DataFrame(data_updated)
out_updated = df_updated.groupby('ID', group_keys=False).apply(get_flag)
print("\n更新数据输出：")
print(out_updated)

代码说明 #

1. 先定位每组第一个'S'的位置，无则直接返回全0
2. 标记第一个'S'为组号1，后续行按规则判断：
   • 日期差≥30天且为'S'：更新基准与组号，标记新组号
   • 日期差≥30天但非'S'：标记0
   • 日期差<30天：沿用当前组号
3. 第一个'S'之前的所有行统一标记0

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