pandas GROUP BY AND SUM，但在超过一定数量时创建新行

问题描述

我当前有一个数据集，其中我尝试根据列对行进行分组，并对值为整数的列求和。

但是，问题是，一旦总和达到特定阈值，我想创建一个新行

例如，在下面的数据框中，我尝试根据公司名称对行进行分组并合计权重，但是我不希望权重超过100。

输入数据帧：

公司	权重
a	30
b	45
a	27
a	40
b	57
a	57
b	32

输出数据帧：

公司	权重
a	97
a	57
b	89
b	45

我已尝试使用GROUP BY和SUM，但是它无法检测我是否已达到最大值。

我有什么方法可以实现这一点吗？

如有任何帮助，我们将不胜感激！

---

解决方案

我认为这里有必要的循环，因此为了提高性能，使用numba、修改solution from Divakar、按组调用函数GroupBy.transform，然后聚合sum：

from numba import njit

@njit
def make_groups(x, target):
    result = np.empty(len(x),dtype=np.uint64)
    total = 0
    group = 0
    for i,x_i in enumerate(x):
        total += x_i
        if total >= target:
            group += 1
            total = 0
        result[i] = group
    return result

g = df.groupby("Company")["Weight"].transform(lambda x: make_groups(x.to_numpy(), 100))

df1 = (df.groupby(by=["Company", g])
        .sum()
        .reset_index(1, drop=True)
        .sort_values(['Company','Weight'], ascending=[True, False])
        .reset_index())
print (df1)
  Company  Weight
0       a      97
1       a      57
2       b      89
3       b      45