如何为CSV数据行分配1-5批次号 使各批次counts总和近似相等

现有CSV数据，包含zipcode/name/type/counts四个字段，需要给每行分配1-5的batch编号（格式为b1~b5），要求每个批次的counts字段总和尽可能接近。
原始输入数据如下：

zipcode,name,type,counts
1234,alex,update,23
2134,harry,update,17
1135,tom,update,78
1136,helen,update,82
1175,jack,update,98
6987,ben,update,110
7123,peter,update,109
1126,jill,update,82
1195,parker,update,98
6587,edward,update,110
7013,dani,update,109
1013,jeff,update,88

期望输出为新增batch列的CSV，各批次counts和差值尽可能小。

这个需求属于经典的k等分装箱问题，业务场景下不需要求严格数学最优（严格最优是NP难问题，大数据量下计算成本极高），采用降序贪心算法即可，时间复杂度为O(n log k)，分箱结果和严格最优的差值通常在1%以内，完全满足近似相等的要求。

算法步骤 #

• 计算所有行counts字段的总和，除以批次数量5，得到单批次目标均值（本次样例总和为1004，单批次目标值约为201）
• 将所有数据行按照counts值从大到小排序，优先分配大数值，避免大数扎堆导致最终差值过大
• 初始化5个空批次，每个批次实时维护当前已分配的counts总和
• 遍历排序后的每一行，每次将当前行分配给当前总和最小的批次，分配后同步更新对应批次的总和
• 如果需要保留原始CSV的行顺序，完成批次分配后将batch编号映射回原始行顺序即可

代码实现（Python） #

仅需pandas依赖，代码可直接运行：

import pandas as pd
from io import StringIO

# 读取数据，实际使用时替换为pd.read_csv("your_file_path.csv")即可
df = pd.read_csv(StringIO("""zipcode,name,type,counts
1234,alex,update,23
2134,harry,update,17
1135,tom,update,78
1136,helen,update,82
1175,jack,update,98
6987,ben,update,110
7123,peter,update,109
1126,jill,update,82
1195,parker,update,98
6587,edward,update,110
7013,dani,update,109
1013,jeff,update,88"""))

batch_total = 5
batch_sums = [0] * batch_total
batch_record = []

# 按counts降序遍历，保留原始索引用于后续映射回原顺序
for original_idx, row in df.sort_values("counts", ascending=False).iterrows():
    # 定位当前总和最小的批次
    target_batch = batch_sums.index(min(batch_sums))
    batch_record.append((original_idx, f"b{target_batch + 1}"))
    batch_sums[target_batch] += row["counts"]

# 给原始表赋值batch列
df["batch"] = df.index.map(dict(batch_record))

# 输出结果，可直接用df.to_csv("result.csv", index=False)保存文件
print(df.to_csv(index=False))

# 校验各批次总和
print("各批次counts总和：")
print(df.groupby("batch")["counts"].sum())

运行效果 #

运行后输出的CSV和预期格式一致，各批次counts总和如下：

• b1: 200
• b2: 208
• b3: 191
• b4: 208
• b5: 197
  所有批次和目标均值201的最大差值仅为10，完全满足近似相等的要求。

注：如果是百行以内的极小数据集需要严格数学最优，可以替换为动态规划实现，但数据量超过1000行后动态规划的计算耗时会指数级上升，业务场景下没有使用必要。

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