Pandas - 前 n 个最大值 groupby 熊猫数据框

Question

我有一个熊猫数据框，如下所示：

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 对数据框进行分组，并从每个组的列 b 中返回前 N 个最大值，按最大值降序排列。

我做了什么？

让我们考虑一下我只想要前两个最大值，我已经完成了以下操作：

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 函数，但它没有帮助，因为它删除了其他列并只返回 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) 聚合以选择每个组的前 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