python - Pandas group by 和 sum，但在超过一定数量时创建一个新行

Question

我目前有一个数据集，我试图根据列对行进行分组，并对值为整数的列求和。
但是，问题是我想在总和达到某个阈值后创建一个新行
例如，在下面的数据框中，我试图根据公司名称对行进行分组并对权重求和，但是，我不希望权重超过 100。
输入数据框:


公司
重量


一种
30

乙
45

一种
27

一种
40

乙
57

一种
57

乙
32


输出数据帧:


公司
重量


一种
97

一种
57

乙
89

乙
45


我曾尝试使用 group by 和 sum，但是，它无法检测我是否达到了最大数量。
有什么办法可以实现这一目标吗？
任何帮助将不胜感激!

Answer 1

我认为这里有必要的循环，所以为了提高性能是使用 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