如何用Pandas切割日期区间以整合指定日期

问题场景 #

手里有个大型Pandas数据集，包含多个id和多组日期区间，现在要给每个id的现有区间做切割，把指定的日期整合到区间划分里，还要保留数据里30多列其他信息，而且不能用遍历——数据太大，遍历效率太低。

可复现示例数据 #

import pandas as pd

# 原数据集
have = pd.DataFrame({
    'id': [1, 1, 2],
    'start_dt': pd.to_datetime(['2023-01-01', '2023-03-01', '2023-02-01']),
    'end_dt': pd.to_datetime(['2023-02-28', '2023-04-30', '2023-04-30']),
    'other_col1': ['A', 'B', 'C'],
    'other_col2': [10, 20, 30]
})

# 要整合的指定日期列表
specific_dt = pd.to_datetime(['2023-01-15', '2023-03-15'])

# 期望的输出结果
want = pd.DataFrame({
    'id': [1, 1, 1, 1, 2, 2, 2],
    'start_dt': pd.to_datetime(['2023-01-01', '2023-01-15', '2023-03-01', '2023-03-15', '2023-02-01', '2023-03-15', '2023-03-15']),
    'end_dt': pd.to_datetime(['2023-01-14', '2023-02-28', '2023-03-14', '2023-04-30', '2023-03-14', '2023-03-14', '2023-04-30']),
    'other_col1': ['A', 'A', 'B', 'B', 'C', 'C', 'C'],
    'other_col2': [10, 10, 20, 20, 30, 30, 30]
})

高效解决方案 #

核心就是用Pandas的分组矢量化操作，别搞逐行遍历，步骤如下：

1. 给每个id收集所有切割边界：把指定日期和每个id原区间的起止日期合并，只留下落在原区间内的指定日期作为切割点。
2. 生成新的日期区间：把所有边界日期排序，生成新的起止日期对，同时自动带好原有的其他列。

具体代码：

def split_intervals(group, specific_dates):
    # 拿当前组的原区间起止日期
    starts = group['start_dt'].tolist()
    ends = group['end_dt'].tolist()
    # 提取当前id对应的其他列数据（同id下其他列值一致，直接取第一行）
    other_cols = group.drop(['start_dt', 'end_dt'], axis=1).iloc[0].to_dict()
    
    # 收集所有要用到的边界日期：原区间起止 + 落在区间内的指定日期
    all_dates = []
    for s, e in zip(starts, ends):
        # 筛出当前区间里的指定日期
        in_range_dates = specific_dates[(specific_dates >= s) & (specific_dates <= e)]
        # 把原start、筛选后的指定日期、原end都加进去
        all_dates.extend([s] + in_range_dates.tolist() + [e])
    
    # 去重后排序
    all_dates = sorted(list(set(all_dates)))
    
    # 生成新的区间
    new_intervals = []
    for i in range(len(all_dates)-1):
        new_start = all_dates[i]
        # 中间区间的end是下一个start减1天，最后一个区间保留原end
        new_end = all_dates[i+1] - pd.Timedelta(days=1) if i != len(all_dates)-2 else all_dates[i+1]
        new_intervals.append({
            'start_dt': new_start,
            'end_dt': new_end,
            **other_cols
        })
    
    return pd.DataFrame(new_intervals)

# 按id分组处理，最后重置索引
result = have.groupby('id').apply(split_intervals, specific_dates=specific_dt).reset_index(drop=True)

# 查看结果
print(result)

方法优势 #

• 用groupby批量处理每个id，完全避免逐行循环，大数据量下效率拉满
• 自动保留所有其他列，不用手动写30多列的名字
• 切割逻辑完全匹配示例：指定日期把原区间拆成[原start, 指定日期-1]、[指定日期, 下一个指定日期-1]、[最后一个指定日期, 原end]的结构

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