Pandas 数据框通过聚合数据删除行

Question

我有一个这样的数据框

test1 = pd.DataFrame(np.array([[1, 9, 3], [1, 5, 6], [2, 1, 9]]),
                   columns=['a', 'b', 'c'])

	a	b	c
0	1	9	3
1	1	5	6
2	2	1	9

如果同一 'a' 下的 'b' 之和大于 10，我想保留 'a'。

对于这种情况，期望的输出是：

	a	b	c
0	1	9	3
1	1	5	6

我的解决方案如下：

test1 = pd.DataFrame(np.array([[1, 9, 3], [1, 5, 6], [2, 1, 9]]),
                   columns=['a', 'b', 'c'])
tmp_ = test1.groupby("a").sum().reset_index()
test1[test1["a"].isin(tmp_[tmp_["b"]>10]["a"].to_list())]

我只是想知道是否有更优雅的方式来做到这一点？

Answer 1

按“a”分组并使用transform

test1 = pd.DataFrame(np.array([[1, 9, 3], [1, 5, 6], [2, 1, 9]]),
                     columns=['a', 'b', 'c'])

b_sum_by_a = test1.groupby('a')['b'].transform('sum') 

test1 = test1[b_sum_by_a > 10] 

>>> test1 

   a  b  c
0  1  9  3
1  1  5  6