python - 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，并从每组的列中返回前 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

输出 Works fine，但这不是执行此操作的有效方法。

我也尝试过nlargest函数，但它没有帮助，因为它会删除其他列并仅返回该b列。

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

这给了我

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

做这个的最好方式是什么？

标签： pythonpandasdataframemaxpandas-groupby

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