如何基于时间戳范围分组DataFrame并最大化分组平均规模？

原始DataFrame #

prod_id                    timestamp1                 timestamp2
1                          2023-12-02                 2023-12-01    
2                          2023-12-05                 2023-12-01    
3                          2023-12-06                 2023-12-01    
4                          2023-12-07                 2023-12-01    
5                          2023-12-08                 2023-12-01
6                          2023-12-08                 2023-12-02
7                          2023-10-10                 2023-09-02
8                          2023-12-11                 2023-12-22 
9                          2023-12-12                 2023-12-24

分组需求 #

需要按时间戳范围对prod_id分组（新增group_id字段），满足以下条件：

• 每个分组内timestamp1的最大日期与最小日期之差不超过3天
• 每个分组内timestamp2的最大日期与最小日期之差同样不超过3天
• 同时要最大化每个分组的prod_id平均数量

期望结果 #

prod_id   timestamp1      timestamp2          group_id
1          2023-12-02    2023-12-01               1
2          2023-12-05    2023-12-01               2
3          2023-12-06    2023-12-01               2
4          2023-12-07    2023-12-01               2
5          2023-12-08    2023-12-01               2
6          2023-12-08    2023-12-02               2
7          2023-10-10    2023-09-02               3
8          2023-12-11    2023-12-22               4
9          2023-12-12    2023-12-24               4

示例说明：prod_id 1单独分到group_id 1，prod_id 2-6分到group_id 2。尽管prod_id 2理论上可以加入group_id 1，但当前分组方式能最大化平均分组规模。

现有代码问题 #

我写了以下代码，但它没有实现最大化分组平均规模的逻辑：

df['timestamp1'] = pd.to_datetime(df['timestamp1'])
df['timestamp2'] = pd.to_datetime(df['timestamp2'])
df_sorted = df.sort_values(by=['timestamp1', 'timestamp2'])

group_id = 1
df_sorted['group_id'] = 0

while not df_sorted.empty:
    group_start = df_sorted.iloc[0]
    valid_rows_t1 = df_sorted[(df_sorted['timestamp1'] <= group_start['timestamp1'] + pd.Timedelta(days=30))]
    valid_rows_t2 = valid_rows_t1[(valid_rows_t1['timestamp2'] >= valid_rows_t1['timestamp2'].min()) &
                                  (valid_rows_t1['timestamp2'] <= valid_rows_t1['timestamp2'].min() + pd.Timedelta(days=30))]
    df_sorted.loc[valid_rows_t2.index, 'group_id'] = group_id
    df_sorted = df_sorted.drop(valid_rows_t2.index)
    group_id += 1

创建DataFrame的代码 #

# 创建DataFrame
data = {
    'prod_id': [1, 2, 3, 4, 5, 6, 7, 8, 9],
    'timestamp1': ['2023-12-02', '2023-12-05', '2023-12-06', '2023-12-07', '2023-12-08', '2023-12-08', '2023-10-10', '2023-12-11', '2023-12-12'],
    'timestamp2': ['2023-12-01', '2023-12-01', '2023-12-01', '2023-12-01', '2023-12-01', '2023-12-02', '2023-09-02', '2023-12-22', '2023-12-24']
}
df = pd.DataFrame(data)

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