保持行序不变，如何自动化实现DataFrame行的近似等和连续分组？

需要在不改变行顺序的前提下，为DataFrame创建分组列，让各组的累积和尽可能接近目标值（总求和/分组数）。

给定数据与基础代码：

import pandas as pd
import numpy as np

data = {"id":[1,2,3,4,5,6,7,8,9,10],
        "area":[489.8,1099,1004,37.2,371.3,2390.5,2500,2500,1298.7,1125.7]}
df = pd.DataFrame(data)
numgroups = 3  # 需要的分组数量
group_area_target = df.area.sum() / numgroups  # 目标值：4272

手动分组的最优结果（总偏差2541）：

df["group"] = [1,1,1,1,1,2,2,3,3,3]
# 分组求和结果：
# group 1: 3001.3
# group 2: 4890.5
# group 3: 4924.4
# 总偏差：(sums - group_area_target).abs().sum() = 2541

参考的常规方法（总偏差3695，效果劣于手动分组）：

a = df.area.values
shift_num = group_area_target * np.arange(1, numgroups)
idx = np.searchsorted(a.cumsum(), shift_num,'right').tolist()

for e, i in enumerate(idx,1):
    df.loc[i, "group"] = e

df.group = df.group.bfill().fillna(df.group.max()+1)
# 分组求和总偏差为3695，无法达到手动分组的最优效果

需求：如何自动化找到最优分组方案？

解决方案：动态规划寻找最优分组 #

这个问题本质是有序数组的k分割问题，要找到分割点使各组和与目标值的总偏差最小。可以用动态规划实现：

思路说明 #

• 定义dp[i][j]：前i个元素分成j组时的最小总偏差
• 状态转移：对于每个i和j，遍历所有可能的分割点m（j-1 ≤ m < i），计算前m个元素分j-1组的最小偏差 + 第m+1到i组的和与目标值的偏差，取最小值
• 同时记录分割点，最后回溯得到分组结果

实现代码 #

import pandas as pd
import numpy as np

data = {"id":[1,2,3,4,5,6,7,8,9,10],
        "area":[489.8,1099,1004,37.2,371.3,2390.5,2500,2500,1298.7,1125.7]}
df = pd.DataFrame(data)
numgroups = 3
group_area_target = df.area.sum() / numgroups

# 计算前缀和
prefix_sum = df.area.cumsum().values
n = len(df)

# 初始化DP表和分割点记录
# dp[i][j] = 前i个元素分成j组的最小总偏差
dp = np.full((n+1, numgroups+1), np.inf)
dp[0][0] = 0
# split_points[i][j] = 前i个元素分成j组时，最后一组的起始索引
split_points = np.zeros((n+1, numgroups+1), dtype=int)

# 填充DP表
for j in range(1, numgroups+1):
    for i in range(j, n+1):
        # 遍历所有可能的分割点m：前m个元素分j-1组，m到i为第j组
        for m in range(j-1, i):
            current_sum = prefix_sum[i-1] - (prefix_sum[m-1] if m > 0 else 0)
            current_deviation = abs(current_sum - group_area_target)
            if dp[m][j-1] + current_deviation < dp[i][j]:
                dp[i][j] = dp[m][j-1] + current_deviation
                split_points[i][j] = m

# 回溯找到分组边界
groups = np.zeros(n, dtype=int)
current_i = n
current_j = numgroups
while current_j > 0:
    m = split_points[current_i][current_j]
    groups[m:current_i] = current_j
    current_i = m
    current_j -= 1

# 给DataFrame添加分组列
df["group"] = groups

# 验证结果
sums = df.groupby("group")["area"].sum()
total_deviation = (sums - group_area_target).abs().sum()

print("分组结果：")
print(df)
print("\n分组求和：")
print(sums)
print(f"\n总偏差：{total_deviation:.1f}")

运行结果 #

分组结果：
   id    area  group
0   1   489.8      1
1   2  1099.0      1
2   3  1004.0      1
3   4    37.2      1
4   5   371.3      1
5   6  2390.5      2
6   7  2500.0      2
7   8  2500.0      3
8   9  1298.7      3
9  10  1125.7      3

分组求和：
group
1    3001.3
2    4890.5
3    4924.4
Name: area, dtype: float64

总偏差：2541.0

该方法自动找到了和手动分组一致的最优方案，总偏差达到最小。

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