python - Pandas 损失最大一栏

我有这个数据集：

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

最佳答案

使用缺少的列的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