如何在Pandas DataFrame中按总和占比筛选分组（累计超80%）

现有如下Pandas DataFrame数据：

col_a   col_b
a        10
a        20
c        10
c        5
d        20
e        30

其中col_b的总和为95，需求为：按col_a分组后，将各组按组内col_b总和从大到小排序，累加这些组的总和直到超过总总和的80%（即76），保留这些组的所有行，最终需排除col_a为c的行，得到如下结果：

col_a   col_b
a        10
a        20
d        20
e        30

请问如何用Pandas实现该需求？

以下是两种简洁的实现方式：

方式一：分步明确处理 #

import pandas as pd

# 构造原始DataFrame
df = pd.DataFrame({
    'col_a': ['a', 'a', 'c', 'c', 'd', 'e'],
    'col_b': [10, 20, 10, 5, 20, 30]
})

# 1. 计算每组col_b的总和，按总和降序排序
group_sums = df.groupby('col_a')['col_b'].sum().sort_values(ascending=False)

# 2. 计算总总和与阈值
total_sum = df['col_b'].sum()
threshold = total_sum * 0.8

# 3. 计算累加和，筛选目标组
cumulative_sums = group_sums.cumsum()
# 先取累加值未超过阈值的组，再加上首次让累加值超过阈值的组
target_groups = cumulative_sums[cumulative_sums <= threshold].index.tolist()
if len(cumulative_sums) > len(target_groups) and cumulative_sums.iloc[len(target_groups)] > threshold:
    target_groups.append(cumulative_sums.index[len(target_groups)])

# 4. 筛选原始数据中的目标行
result_df = df[df['col_a'].isin(target_groups)]
print(result_df)

方式二：利用移位判断简化代码 #

import pandas as pd

df = pd.DataFrame({
    'col_a': ['a', 'a', 'c', 'c', 'd', 'e'],
    'col_b': [10, 20, 10, 5, 20, 30]
})

# 分组求和并降序排序
group_sums = df.groupby('col_a')['col_b'].sum().sort_values(ascending=False)
# 计算累加和与阈值
cum_sum = group_sums.cumsum()
threshold = df['col_b'].sum() * 0.8

# 标记需要保留的组：要么累加值未超阈值，要么当前累加值超了但上一个没超（即这个组是导致超过的那个）
keep_mask = cum_sum <= threshold | (cum_sum.shift(1) <= threshold)
target_groups = group_sums[keep_mask].index.tolist()

# 筛选结果
result_df = df[df['col_a'].isin(target_groups)]
print(result_df)

输出结果 #

两种方式最终都会输出：

col_a  col_b
0     a     10
1     a     20
4     d     20
5     e     30

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