如何用Pandas拆分包含指定日期的日期区间行

问题描述 #

现有如下Pandas DataFrame，包含分组、日期区间和对应数值：

import pandas as pd

data = [['A', '2022-09-01', '2022-09-05', 10], ['A', '2022-09-05', '2022-09-15', 1], ['A', '2022-09-15', '2022-09-18', 12], ['B', '2022-09-01', '2022-09-03', 4], ['B', '2022-09-03', '2022-09-07', 7], ['B', '2022-09-07', '2022-09-12', 9]]
df = pd.DataFrame(data, columns=['GROUP', 'start_date', 'end_date', 'value'])

需要针对指定日期（如2022-09-10），按组检查每一行：若指定日期落在该行的start_date与end_date区间内，则拆分该行：

• 原行的end_date替换为指定日期
• 新增一行，start_date为指定日期，end_date为原行的end_date，两行的value保持不变

期望输出示例：

certain_date = '2022-09-10'

data = [['A', '2022-09-01', '2022-09-05', 10], ['A', '2022-09-05', '2022-09-10', 1], ['A', '2022-09-10', '2022-09-15', 1], ['A', '2022-09-15', '2022-09-18', 12], ['B', '2022-09-01', '2022-09-03', 4], ['B', '2022-09-03', '2022-09-07', 7], ['B', '2022-09-07', '2022-09-10', 9], ['B', '2022-09-10', '2022-09-12', 9]]
df_desired = pd.DataFrame(data, columns=['GROUP', 'start_date', 'end_date', 'value'])

实现步骤 #

1. 转换日期列类型 #

先将start_date和end_date转换为datetime类型，确保日期比较的准确性：

df['start_date'] = pd.to_datetime(df['start_date'])
df['end_date'] = pd.to_datetime(df['end_date'])
certain_date = pd.to_datetime(certain_date)

2. 筛选需要拆分的行 #

找出满足start_date < certain_date < end_date的目标行：

mask = (df['start_date'] < certain_date) & (df['end_date'] > certain_date)
to_split = df[mask].copy()

3. 生成拆分后的新行 #

对筛选出的行，生成拆分后的两组数据：

# 原行修改end_date为指定日期
split_row1 = to_split.assign(end_date=certain_date)
# 新行修改start_date为指定日期，保留原end_date
split_row2 = to_split.assign(start_date=certain_date)

4. 合并数据并排序 #

将不需要拆分的行、拆分后的两行合并，再按GROUP和start_date排序，恢复原始顺序：

# 提取不需要拆分的行
no_split = df[~mask]
# 合并所有数据
result = pd.concat([no_split, split_row1, split_row2], ignore_index=True)
# 按分组和起始日期排序，保证顺序正确
result = result.sort_values(by=['GROUP', 'start_date'], ignore_index=True)

完整代码 #

import pandas as pd

# 原始数据
data = [['A', '2022-09-01', '2022-09-05', 10], ['A', '2022-09-05', '2022-09-15', 1], ['A', '2022-09-15', '2022-09-18', 12], ['B', '2022-09-01', '2022-09-03', 4], ['B', '2022-09-03', '2022-09-07', 7], ['B', '2022-09-07', '2022-09-12', 9]]
df = pd.DataFrame(data, columns=['GROUP', 'start_date', 'end_date', 'value'])

# 指定日期
certain_date = '2022-09-10'

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

# 筛选需要拆分的行
mask = (df['start_date'] < certain_date) & (df['end_date'] > certain_date)
to_split = df[mask].copy()

# 生成拆分后的行
split_row1 = to_split.assign(end_date=certain_date)
split_row2 = to_split.assign(start_date=certain_date)

# 合并并排序
no_split = df[~mask]
result = pd.concat([no_split, split_row1, split_row2], ignore_index=True)
result = result.sort_values(by=['GROUP', 'start_date'], ignore_index=True)

# 可选：转换回字符串格式（如果需要）
result['start_date'] = result['start_date'].dt.strftime('%Y-%m-%d')
result['end_date'] = result['end_date'].dt.strftime('%Y-%m-%d')

print(result)

运行后即可得到符合需求的DataFrame。

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