Pandas - 前 n 个最大值 groupby pandas 数据帧

Question

我有一个 pandas 数据框如下：

df = pd.DataFrame({
    'a': [1, 1, 1, 1, 1, 2, 2, 2, 2],
    'b': [3, 2, 1, 4, 2, 1, 2, 2, 1]
})

这给了我

>>> df
    a   b
0   1   3
1   1   2
2   1   1
3   1   4
4   1   2
5   2   1
6   2   2
7   2   2
8   2   1

我想按列a对数据帧进行分组，return前N个最大值来自每组的第 b 列，按最大值降序排列。

我做了什么？

假设我只想要前两个最大值，我做了以下操作：

1. 按两列降序对数据框进行排序
2. 获取前两个值
3. 由于这些值是升序排列的，所以再按降序排列

df = df.sort_values(['a', 'b'], ascending=False)
df = df.groupby('a').nth([0,1])
df = df.sort_values(['a', 'b'], ascending=False).reset_index()

这给了我

>>> df
    a   b
0   2   2
1   2   2
2   1   4
3   1   3

输出工作正常，但这不是执行此操作的有效方法。

我也试过 nlargest 函数，但没有用，因为它会删除其他列，return 只会删除 b 列。

df = df.groupby('a')['b'].nlargest(2)

这给了我

>>> df
0    2
1    2
2    4
3    3
Name: b, dtype: int64

最好的方法是什么？

Answer 1

Sort 按列的值 b 然后 group 数据框并使用 head(n) 聚合到 select 的前 n 行每组

df.sort_values('b', ascending=False).groupby('a').head(2)

---

   a  b
3  1  4
0  1  3
6  2  2
7  2  2