熊猫在每组中获得最高n条记录

Question

假设我有像这样的pandas DataFrame：>>> df = pd.DataFrame({'id':[1,1,1,2,2,2,2,3,4],'value':[1,2,3,1,2,3,4,1,1]})>>> df   id  value0   1      11   1      22   1      33   2      14   2      25   2      36   2      47   3      18   4      1我想为每个id获取一个包含前2条记录的新DataFrame，如下所示：   id  value0   1      11   1      23   2      14   2      27   3      18   4      1我可以通过以下方式在组内编号记录：>>> dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()>>> dfN   id  level_1  index  value0   1        0      0      11   1        1      1      22   1        2      2      33   2        0      3      14   2        1      4      25   2        2      5      36   2        3      6      47   3        0      7      18   4        0      8      1>>> dfN[dfN['level_1'] <= 1][['id', 'value']]   id  value0   1      11   1      23   2      14   2      27   3      18   4      1但这样做有更有效/优雅的方法吗？并且每个组中的数字记录还有更优雅的方法（如SQL窗口函数row_number（））。

Answer 1

你试过了吗 df.groupby('id').head(2)

产生的输出：

>>> df.groupby('id').head(2)

       id  value

id             

1  0   1      1

   1   1      2 

2  3   2      1

   4   2      2

3  7   3      1

4  8   4      1

（请记住，您可能需要先订购/排序，具体取决于您的数据）

编辑：正如提问者所提到的，用于df.groupby('id').head(2).reset_index(drop=True)删除多索引并展平结果。

>>> df.groupby('id').head(2).reset_index(drop=True)

    id  value

0   1      1

1   1      2

2   2      1

3   2      2

4   3      1

5   4      1

Answer 2

由于0.14.1，你现在可以做的nlargest和nsmallest一个上groupby对象：

In [23]: df.groupby('id')['value'].nlargest(2)

Out[23]: 

id   

1   2    3

    1    2

2   6    4

    5    3

3   7    1

4   8    1

dtype: int64

还有，你在那里得到的原始索引以及轻微的怪事，但根据您的原始索引是什么，这可能是真正有用的是。

如果你对它不感兴趣，你可以.reset_index(level=1, drop=True)完全摆脱它。

（注意：从0.17.1开始，你也可以在DataFrameGroupBy上执行此操作，但现在它只适用于Series和SeriesGroupBy。）