pandas如何对DataFrame按列数值分组使每组总和近似目标值

核心思路 #

采用贪心累加算法实现，遍历每行数据时累计当前分组的Quantity总和，当加入当前行数值会超过设定的目标总和时，就开启新分组。该算法时间复杂度为O(n)，对大批量数据非常友好，分组总和误差可控，完全满足近似相等的需求。

代码实现 #

import pandas as pd

def add_quantity_group(df: pd.DataFrame, target_sum: int, sort_first: bool = False) -> pd.DataFrame:
    res_df = df.copy()
    # 可选先按Quantity排序，减少分组误差，如需保留原顺序可忽略该步骤
    if sort_first:
        res_df = res_df.sort_values('Quantity', ascending=False).reset_index(drop=True)
    
    current_total = 0
    group_num = 0
    group_list = []
    
    for qty in res_df['Quantity']:
        if current_total + qty > target_sum:
            group_num += 1
            current_total = qty
        else:
            current_total += qty
        group_list.append(group_num)
    
    res_df['group'] = group_list
    # 如果之前做了排序，恢复原数据顺序
    if sort_first:
        res_df = res_df.sort_index().reset_index(drop=True)
    return res_df

使用示例 #

# 构造测试数据
df = pd.DataFrame({
    'ID': range(1, 11),
    'Quantity': [10, 22, 18, 15, 7, 25, 3, 12, 9, 14]
})

# 场景1：单组目标总和为50
df_group_50 = add_quantity_group(df, target_sum=50)
# 查看分组总和验证效果
print(df_group_50.groupby('group')['Quantity'].sum())
"""
输出结果近似如下：
group
0    50
1    47
2    38
Name: Quantity, dtype: int64
"""

# 场景2：单组目标总和为30
df_group_30 = add_quantity_group(df, target_sum=30)
print(df_group_30.groupby('group')['Quantity'].sum())
"""
输出结果近似如下：
group
0    32
1    15
2    32
3    24
4    14
Name: Quantity, dtype: int64
"""

# 场景3：固定分批数量，每批总和大致相等
total_sum = df['Quantity'].sum()
batch_num = 3 # 要分3批
target = total_sum // batch_num
df_batch = add_quantity_group(df, target_sum=target)

注意事项 #

• 若存在单个Quantity值本身超过目标总和的情况，该行会单独成为一个分组，符合分批处理的实际逻辑
• 开启sort_first=True参数可以让分组总和更接近目标值，但会额外消耗排序的性能开销，超大数据量下可根据需求选择
• 该方法不需要加载全量数据到内存后再计算，也可适配流式分批处理场景

内容的提问来源于stack exchange，提问作者Rob E.