如何按两个时间戳的日期范围对数据集行进行分组？

现有一个包含两个时间戳字段的DataFrame：

timestamp1             timestamp2
2022-02-18            2023-01-02    
2022-02-19            2023-01-04    
2022-02-21            2023-01-11    
2022-03-11            2024-02-05    
2022-03-12            2024-02-06    
2022-03-30            2024-02-07        

需要按时间戳范围对行进行分组（新增group字段），要求：

• 组内每个时间戳字段的日期范围不超过4天（即组内该字段的最大值与最小值差值≤4天）
• 在满足约束的前提下，每个组包含尽可能多的行

期望结果如下：

timestamp1             timestamp2       group
2022-02-18            2023-01-02            1
2022-02-19            2023-01-04            1
2022-02-21            2023-01-11            2
2022-03-11            2024-02-05            3
2022-03-12            2024-02-06            3
2022-03-30            2024-02-07            4

分组规则说明：

• 第三行不属于组1：其timestamp2与组1最小timestamp2差值为9天，超出4天限制
• 最后一行不属于组3：其timestamp1与组3最小timestamp1差值为19天，超出4天限制

尝试的代码无法得到期望结果，尤其在处理大型数据集时：

import pandas as pd

# Sample DataFrame
data = {
    'timestamp1': ['2022-02-18', '2022-02-19', '2022-02-21', '2022-03-11', '2022-03-12', '2022-03-30'],
    'timestamp2': ['2023-01-02', '2023-01-04', '2023-01-11', '2024-02-05', '2024-02-06', '2024-02-07']
}
df = pd.DataFrame(data)
df['timestamp1'] = pd.to_datetime(df['timestamp1'])
df['timestamp2'] = pd.to_datetime(df['timestamp2'])

# Function to assign groups
def assign_groups(df):
    groups = []
    current_group = 1
    start_timestamp2 = df.iloc[0]['timestamp2']
    for i, row in df.iterrows():
        if (row['timestamp2'] - start_timestamp2).days > 4:
            current_group += 1
            start_timestamp2 = row['timestamp2']
        groups.append(current_group)
    return groups

# Assign groups
df['group'] = assign_groups(df)

原代码仅检查了timestamp2与组起始值的差值，未同时验证timestamp1的组内范围，也未考虑组内所有行的最大最小差值，导致分组不符合要求。以下是修正后的实现：

import pandas as pd

# 初始化数据
data = {
    'timestamp1': ['2022-02-18', '2022-02-19', '2022-02-21', '2022-03-11', '2022-03-12', '2022-03-30'],
    'timestamp2': ['2023-01-02', '2023-01-04', '2023-01-11', '2024-02-05', '2024-02-06', '2024-02-07']
}
df = pd.DataFrame(data)
df['timestamp1'] = pd.to_datetime(df['timestamp1'])
df['timestamp2'] = pd.to_datetime(df['timestamp2'])

def assign_groups(df):
    if df.empty:
        return []
    
    groups = []
    current_group = 1
    # 记录当前组的两个时间戳字段的最小、最大值
    group_min_ts1 = df.iloc[0]['timestamp1']
    group_max_ts1 = df.iloc[0]['timestamp1']
    group_min_ts2 = df.iloc[0]['timestamp2']
    group_max_ts2 = df.iloc[0]['timestamp2']
    
    groups.append(current_group)
    
    for i in range(1, len(df)):
        row_ts1 = df.iloc[i]['timestamp1']
        row_ts2 = df.iloc[i]['timestamp2']
        
        # 计算将当前行加入组后的新范围
        new_min_ts1 = min(group_min_ts1, row_ts1)
        new_max_ts1 = max(group_max_ts1, row_ts1)
        new_min_ts2 = min(group_min_ts2, row_ts2)
        new_max_ts2 = max(group_max_ts2, row_ts2)
        
        # 验证两个时间戳的范围是否均≤4天
        ts1_diff = (new_max_ts1 - new_min_ts1).days
        ts2_diff = (new_max_ts2 - new_min_ts2).days
        
        if ts1_diff <= 4 and ts2_diff <= 4:
            # 符合条件，加入当前组并更新组范围
            group_min_ts1 = new_min_ts1
            group_max_ts1 = new_max_ts1
            group_min_ts2 = new_min_ts2
            group_max_ts2 = new_max_ts2
            groups.append(current_group)
        else:
            # 不符合条件，新建组并重置范围
            current_group += 1
            group_min_ts1 = row_ts1
            group_max_ts1 = row_ts1
            group_min_ts2 = row_ts2
            group_max_ts2 = row_ts2
            groups.append(current_group)
    
    return groups

df['group'] = assign_groups(df)
print(df)

代码说明 #

1. 初始化时记录第一个组的两个时间戳字段的范围
2. 遍历后续每一行，模拟将其加入当前组，计算新的时间戳范围
3. 若两个字段的范围差值均≤4天，则加入当前组并更新组范围；否则新建组
4. 该逻辑确保每个组在满足约束的前提下，包含尽可能多的行

运行结果与期望一致：

timestamp1 timestamp2  group
0 2022-02-18 2023-01-02      1
1 2022-02-19 2023-01-04      1
2 2022-02-21 2023-01-11      2
3 2022-03-11 2024-02-05      3
4 2022-03-12 2024-02-06      3
5 2022-03-30 2024-02-07      4

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