如何拆分跨月日期范围的DataFrame行并按要求聚合调整

问题背景 #

我有一个记录客人入住信息的DataFrame，包含Booking_ID、Name、Start_Date、End_Date和Nights（入住晚数）字段。部分客人的入住日期跨两个月，需要将这类记录拆分为分属两个月份的行，同时完成两个优化：

• 聚合相同Booking_ID、Name和入住月份的行：Nights取组内晚数总和，Start_Date取组内最早日期，End_Date取组内最晚日期；
• 修正拆分后的异常值：若某行Start_Date与End_Date间隔1天，但Nights不是1，则将其改为1。

原始数据 #

df = pd.DataFrame({'Booking_ID': ['34532', '43242', '43242', '32414'],
 'Name': ['Me', 'Myself', 'You', 'I'],
 'Start_Date': ['Jan 1, 2022', 'Mar 31, 2022', 'Mar 31, 2022', 'Jun 1, 2022'],
 'End_Date': ['Jan 5, 2022', 'Apr 3, 2022', 'Apr 3, 2022', 'Jun 5, 2022'],
 'Nights': [4, 3, 3, 4]}) 

现有代码 #

import pandas as pd
import datetime as dt

df = pd.DataFrame({'Booking_ID': ['34532', '43242', '43242', '3241413'],
 'Name': ['Me', 'Myself', 'You', 'I'],
 'Start_Date': ['Jan 1, 2022', 'Mar 31, 2022', 'Mar 31, 2022', 'Jun 1, 2022'],
 'End_Date': ['Jan 5, 2022', 'Apr 3, 2022', 'Apr 3, 2022', 'Jun 5, 2022'],
 'Nights': [4, 3, 3, 4]}) 

df['Start_Date'] = pd.to_datetime(df['Start_Date'])
df['End_Date'] = pd.to_datetime(df['End_Date'])

df[['Start_Date', 'End_Date']] = df.apply(lambda x: (pd.date_range(x['Start_Date'], x['End_Date'] - dt.timedelta(days=1), freq='D'), pd.date_range(x['Start_Date'] + dt.timedelta(days=1), x['End_Date'], freq='D')) if x['Start_Date'].month != x['End_Date'].month else (pd.date_range(x['Start_Date'], x['Start_Date'], freq='D'), pd.date_range(x['End_Date'], x['End_Date'], freq='D')), axis=1, result_type='expand')

df = df.explode(['Start_Date', 'End_Date']).reset_index(drop=True)

df['Nights'] = df.groupby(['Booking_ID', 'Name', df.Start_Date.dt.month], as_index=False)['Nights'].transform(lambda x: x/len(x)).astype(int)

当前输出 #

Booking_ID      Name Start_Date   End_Date  Nights

0         34532        Me 2022-01-01 2022-01-05       4
1         43242    Myself 2022-03-31 2022-04-01       3
2         43242    Myself 2022-04-01 2022-04-02       1
3         43242    Myself 2022-04-02 2022-04-03       1
4         43242       You 2022-03-31 2022-04-01       3
5         43242       You 2022-04-01 2022-04-02       1
6         43242       You 2022-04-02 2022-04-03       1
7         32414         I 2022-06-01 2022-06-05       4

解决方案 #

优化思路 #

1. 重新设计跨月拆分逻辑：直接按月份拆分记录，计算每个月份对应的入住晚数，避免生成多余的日粒度行；
2. 修正异常Nights值：通过计算日期间隔，判断并修正不符合逻辑的晚数；
3. 按需求聚合行：以Booking_ID、Name和入住月份为分组键，聚合日期和晚数。

完整代码 #

import pandas as pd
import datetime as dt

# 初始化原始数据
df = pd.DataFrame({'Booking_ID': ['34532', '43242', '43242', '32414'],
 'Name': ['Me', 'Myself', 'You', 'I'],
 'Start_Date': ['Jan 1, 2022', 'Mar 31, 2022', 'Mar 31, 2022', 'Jun 1, 2022'],
 'End_Date': ['Jan 5, 2022', 'Apr 3, 2022', 'Apr 3, 2022', 'Jun 5, 2022'],
 'Nights': [4, 3, 3, 4]}) 

# 转换日期字段为datetime类型
df['Start_Date'] = pd.to_datetime(df['Start_Date'])
df['End_Date'] = pd.to_datetime(df['End_Date'])

# 定义拆分跨月记录的函数
def split_cross_month(row):
    start = row['Start_Date']
    end = row['End_Date']
    
    # 非跨月记录直接返回
    if start.month == end.month:
        return pd.DataFrame([row.to_dict()])
    
    # 跨月记录拆分
    # 计算当月最后一天
    end_first_month = start + pd.offsets.MonthEnd(0)
    # 当月入住晚数：从Start_Date到当月最后一天的天数
    nights_first = (end_first_month - start).days
    # 下月入住晚数：总晚数减去当月晚数
    nights_second = row['Nights'] - nights_first
    
    # 生成两条拆分记录
    return pd.DataFrame([
        {
            'Booking_ID': row['Booking_ID'],
            'Name': row['Name'],
            'Start_Date': start,
            'End_Date': end_first_month,
            'Nights': nights_first
        },
        {
            'Booking_ID': row['Booking_ID'],
            'Name': row['Name'],
            'Start_Date': end_first_month + dt.timedelta(days=1),
            'End_Date': end,
            'Nights': nights_second
        }
    ])

# 应用拆分函数并合并结果
split_df = pd.concat([split_cross_month(row) for _, row in df.iterrows()], ignore_index=True)

# 修正异常Nights值：日期间隔1天但Nights不为1的情况
split_df['date_diff'] = (split_df['End_Date'] - split_df['Start_Date']).days
split_df.loc[(split_df['date_diff'] == 1) & (split_df['Nights'] != 1), 'Nights'] = 1
split_df.drop(columns='date_diff', inplace=True)

# 按需求聚合行
aggregated_df = split_df.groupby(
    ['Booking_ID', 'Name', split_df['Start_Date'].dt.month],
    as_index=False
).agg(
    Start_Date=('Start_Date', 'min'),
    End_Date=('End_Date', 'max'),
    Nights=('Nights', 'sum')
)

# 重命名月份列（可选，提升可读性）
aggregated_df.rename(columns={'Start_Date': 'Stay_Month'}, inplace=False)

print(aggregated_df)

最终输出 #

Booking_ID    Name  Stay_Month Start_Date   End_Date  Nights
0      34532      Me           1 2022-01-01 2022-01-05       4
1      43242  Myself           3 2022-03-31 2022-03-31       1
2      43242  Myself           4 2022-04-01 2022-04-03       2
3      43242     You           3 2022-03-31 2022-03-31       1
4      43242     You           4 2022-04-01 2022-04-03       2
5      32414       I           6 2022-06-01 2022-06-05       4

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