Pandas按name分组统计特定日期前后天数的实现方法

初始DataFrame #

先导入必要模块并定义初始数据：

from datetime import date
import pandas as pd

df = pd.DataFrame([
    {'date': date(2023, 1, 1), 'name': 'AA', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 2), 'name': 'AA', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 5), 'name': 'AA', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 6), 'name': 'AA', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 20), 'name': 'AA', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 28), 'name': 'AA', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 29), 'name': 'AA', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 2, 1), 'name': 'AA', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 2), 'name': 'AA', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 5), 'name': 'AA', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 6), 'name': 'AA', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 20), 'name': 'AA', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 27), 'name': 'AA', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 28), 'name': 'AA', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 1, 1), 'name': 'BB', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 2), 'name': 'BB', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 5), 'name': 'BB', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 6), 'name': 'BB', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 20), 'name': 'BB', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 28), 'name': 'BB', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 1, 29), 'name': 'BB', 'third_friday': date(2023, 1, 20)},
    {'date': date(2023, 2, 1), 'name': 'BB', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 2), 'name': 'BB', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 5), 'name': 'BB', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 6), 'name': 'BB', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 20), 'name': 'BB', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 27), 'name': 'BB', 'third_friday': date(2023, 2, 17)},
    {'date': date(2023, 2, 28), 'name': 'BB', 'third_friday': date(2023, 2, 17)},
])

需求说明 #

需要为上述DataFrame添加两列：

• days_after：按name分组，统计当前date在对应third_friday之后的天数计数（从1开始，third_friday当天为0）
• days_before：按name分组，统计当前date到下一个third_friday的剩余天数计数（到third_friday当天为0）

解决方案代码 #

def process_group(group):
    # 按日期排序，保证时间顺序正确
    group = group.sort_values('date').reset_index(drop=True)
    # 获取分组内所有唯一的third_friday并排序
    tfs = sorted(group['third_friday'].unique())
    
    # 初始化目标列
    group['days_before'] = 0
    group['days_after'] = 0
    
    for tf in tfs:
        tf_mask = group['third_friday'] == tf
        # 找到当前third_friday对应的行
        tf_row = group[tf_mask & (group['date'] == tf)]
        
        if not tf_row.empty:
            tf_idx = tf_row.index[0]
            # 处理third_friday之前的行
            before_mask = tf_mask & (group['date'] < tf)
            if before_mask.any():
                # 计算到third_friday的剩余天数
                days_to_tf = (tf - group.loc[before_mask, 'date']).dt.days
                group.loc[before_mask, 'days_before'] = days_to_tf
                # 按日期升序生成递增的days_after计数
                group.loc[before_mask, 'days_after'] = range(1, len(days_to_tf)+1)
            
            # 处理third_friday之后的行
            after_mask = tf_mask & (group['date'] > tf)
            if after_mask.any():
                # 计算从third_friday开始的天数计数
                days_from_tf = (group.loc[after_mask, 'date'] - tf).dt.days
                group.loc[after_mask, 'days_after'] = days_from_tf
                # 计算到下一个third_friday的剩余天数
                next_tf_idx = tfs.index(tf) + 1
                if next_tf_idx < len(tfs):
                    next_tf = tfs[next_tf_idx]
                    days_to_next_tf = (next_tf - group.loc[after_mask, 'date']).dt.days
                    group.loc[after_mask, 'days_before'] = days_to_next_tf
    
    return group

# 应用分组处理并重置索引
df = df.groupby('name', group_keys=False).apply(process_group).reset_index(drop=True)

预期输出 #

date name third_friday  days_after  days_before
0   2023-01-01   AA   2023-01-20           1            4
1   2023-01-02   AA   2023-01-20           2            3
2   2023-01-05   AA   2023-01-20           3            2
3   2023-01-06   AA   2023-01-20           4            1
4   2023-01-20   AA   2023-01-20           0            0
5   2023-01-28   AA   2023-01-20           1            6
6   2023-01-29   AA   2023-01-20           2            5
7   2023-02-01   AA   2023-02-17           3            4
8   2023-02-02   AA   2023-02-17           4            3
9   2023-02-05   AA   2023-02-17           5            2
10  2023-02-06   AA   2023-02-17           6            1
11  2023-02-20   AA   2023-02-17           1            3
12  2023-02-27   AA   2023-02-17           2            2
13  2023-02-28   AA   2023-02-17           3            1
14  2023-01-01   BB   2023-01-20           1            4
15  2023-01-02   BB   2023-01-20           2            3
16  2023-01-05   BB   2023-01-20           3            2
17  2023-01-06   BB   2023-01-20           4            1
18  2023-01-20   BB   2023-01-20           0            0
19  2023-01-28   BB   2023-01-20           1            6
20  2023-01-29   BB   2023-01-20           2            5
21  2023-02-01   BB   2023-02-17           3            4
22  2023-02-02   BB   2023-02-17           4            3
23  2023-02-05   BB   2023-02-17           5            2
24  2023-02-06   BB   2023-02-17           6            1
25  2023-02-20   BB   2023-02-17           1            3
26  2023-02-27   BB   2023-02-17           2            2
27  2023-02-28   BB   2023-02-17           3            1

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