pandas - Pandas 数据框中 groupby.agg 的 N 最大/idxmax

Question

我想查找的不是单个最大索引，而是 pandas groupby.agg 操作的 n 个最大索引。

我的数据框如下所示:

>>> dat = [{"date": datetime.date(2021,1,1), "category": "cats", "var": 1},
       {"date": datetime.date(2021,1,1), "category": "dogs", "var": 0},
       {"date": datetime.date(2021,2,1), "category": "cats", "var": -1},
       {"date": datetime.date(2021,2,1), "category": "dogs", "var": 3},
       ]
>>> df = pd.DataFrame(dat)

>>> df.groupby(["date", "category"]).agg({'var': ['sum']})
                    var
                    sum
date       category    
2021-01-01 cats       1
           dogs       0
2021-02-01 cats      -1
           dogs       3

获取var值最大的索引:

>>> df.groupby(["date", "category"]).agg({'var': ['sum']}).idxmax()

var  sum    (2021-02-01, dogs)
dtype: object

确认，没错。我可以看到索引 (2021-02-01, dogs) 其中 var 的值为 a a max。

所以现在，我不仅要查看最大索引，还要查看 nlargest 索引。这里，n=2。

然而，对此的明显解决方案 nlargest(2, variable_name) 似乎失败了:

df.groupby(["date", "category"]).agg({'var': ['sum']}).nlargest(2, 'sum')

失败，似乎是多索引失败:

pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()
pandas/_libs/index.pyx in pandas._libs.index.IndexEngine.get_loc()
pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()
pandas/_libs/hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()
KeyError: 'sum'

如何获取 n 个最大索引的 idxmax？我假设有一个多索引可以传递给 nlargest 或 level，但我没有看到它。

Answer 1

怎么样:

(df.groupby(['date','category'])['var'].sum()
   .nlargest(2).index
)

输出:

MultiIndex([('2021-1-1', 'dogs'),
            ('2021-1-1', 'cats')],
           names=['date', 'category'])