如何在pandas中选择每组的前3行？

Question

我得到一个像这样的熊猫数据框：

    id   prob
0    1   0.5   
1    1   0.6
2    1   0.4
3    1   0.2
4    2   0.3
6    2   0.5
...

我想按“id”对其进行分组，按降序排序并获得每组的前 3 个概率。请注意，某些组包含少于 3 的行。最后我想得到一个二维数组，如：

[[1, 0.6, 0.5, 0.4], [2, [0.5, 0.3]]...]

我怎么能用熊猫做到这一点？谢谢！

Answer 1

使用sort_values、groupby、 和head：

df.sort_values(by=['id','prob'], ascending=[True,False]).groupby('id').head(3).values

输出：

array([[ 1. ,  0.6],
       [ 1. ,  0.5],
       [ 1. ,  0.4],
       [ 2. ,  0.5],
       [ 2. ,  0.3]])

关注@COLDSPEED 线索：

df.sort_values(by=['id','prob'], ascending=[True,False])\
  .groupby('id').agg(lambda x: x.head(3).tolist())\
  .reset_index().values.tolist()

输出：

[[1, [0.6, 0.5, 0.4]], [2, [0.5, 0.3]]]