如何在 Pandas DataFrame 中找到 2 行总和大于某个值的行？

Question

在像下面这样的数据集中，我试图按 attr_1 和 attr_2 对行进行分组，如果 count 的总和 em> 列超过阈值（在本例中为 100），我想保留原始行。

account	attr_1	attr_2	count
ABC	X1	Y1	25
DEF	X1	Y1	100
ABC	X2	Y2	150
DEF	X2	Y2	0
ABC	X3	Y3	10
DEF	X3	Y3	15

我正在使用下面的混乱方法，但我想看看是否有更清洁的方法可以处理这个问题。

df = pd.DataFrame({'account': ['ABC', 'DEF','ABC', 'DEF','ABC', 'DEF'],
                   'attr_1': ['X1', 'X1', 'X2', 'X2', 'X3', 'X3'],
                   'attr_2': ['Y1', 'Y1', 'Y2', 'Y2', 'Y3', 'Y3'],
                   'count': [25, 100, 150, 0, 10, 15]
                  })

min_count = 100
groups = df.groupby(by=['attr_1', 'attr_2']).sum()
group_count = groups.apply(lambda g: g[g >= min_count])

# find indices of groups exceed the threshold
keep_index = []
for ix in group_count.index:
    keep_index.extend(df.query(f'attr_1=="{ix[0]}" & attr_2=="{ix[1]}"').index.values)
    
# filter dataframe
output_df = df[df.index.isin(keep_index)]

Answer 1

您可以使用groupby + filter，并在filter lambda 中为组提供标量条件：

df.groupby(['attr_1', 'attr_2']).filter(lambda g:  g['count'].sum() >= min_count)

  account attr_1 attr_2  count
0     ABC     X1     Y1     25
1     DEF     X1     Y1    100
2     ABC     X2     Y2    150
3     DEF     X2     Y2      0

或者使用groupby + transform 创建一个与原始数据框兼容的过滤条件：

df[df.groupby(['attr_1', 'attr_2'])['count'].transform('sum').ge(min_count)]

  account attr_1 attr_2  count
0     ABC     X1     Y1     25
1     DEF     X1     Y1    100
2     ABC     X2     Y2    150
3     DEF     X2     Y2      0