python - pandas nlargest lost one 专栏

Question

我有这个数据集:

Id   query  count
001  abc    20
001  bcd    30
001  ccd   100
002  ace   13
002  ahhd   30
002  ahe    28

我想根据计数找到每个 Id 的 Top2 查询。所以我想看看:

Id   query  count
001  ccd    100
001  bcd    30
002  ahhd   30
002  ahe    28

我试过这两行代码:

df.groupby('Id')['count'].nlargest(2)，结果中丢失了“查询”列，这不是我想要的。那么如何在我的结果中保留查询。编号

001     100
001     30
002     30
002     28

Answer 1

使用set_index缺失列数:

df = df.set_index('query').groupby('Id')['count'].nlargest(2).reset_index()
print (df)
    Id query  count
0  001   ccd    100
1  001   bcd     30
2  002  ahhd     30
3  002   ahe     28