如何使用列值作为pandas聚合函数中的参数？

Question

给定一个特定类型的表

A	乙	C
t	r	1
t	r	1
n	j	2
n	j	2
n	j	2

我想对Aand进行分组B，并且仅采用由指定的行数C

所以期望的输出是

A	乙	C
t	r	1
n	j	2
n	j	2

我试图通过这个功能来实现这一点，但没有运气

df.groupby(['A', 'B']).agg(lambda x: x.head(df.C))

Answer 1

您可以使用groupby.cumcount布尔索引：

out = df[df['C'].gt(df.groupby(['A', 'B']).cumcount())]

或者用经典的groupby.apply：

(df.groupby(['A', 'B'], sort=False, as_index=False, group_keys=False)
   .apply(lambda g: g.head(g['C'].iloc[0]))
)

输出：

   A  B  C
0  t  r  1
2  n  j  2
3  n  j  2

该方法的中间体groupby.cumcount：

   A  B  C  cumcount  C > cumcount
0  t  r  1         0          True
1  t  r  1         1         False
2  n  j  2         0          True
3  n  j  2         1          True
4  n  j  2         2         False