假设我有这样的熊猫数据框架:

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

它看起来像:

   id  value
0   1      1
1   1      2
2   1      3
3   2      1
4   2      2
5   2      3
6   2      4
7   3      1
8   4      1

我想获得一个新的DataFrame与前2记录为每个id,像这样:

   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

我可以用编号记录在组后groupby:

dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()

它看起来像:

   id  level_1  index  value
0   1        0      0      1
1   1        1      1      2
2   1        2      2      3
3   2        0      3      1
4   2        1      4      2
5   2        2      5      3
6   2        3      6      4
7   3        0      7      1
8   4        0      8      1

然后对于期望的输出:

dfN[dfN['level_1'] <= 1][['id', 'value']]

输出:

   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

但是有没有更有效/更优雅的方法来做到这一点呢?此外,每个组中是否有更优雅的方法来处理数字记录(如SQL窗口函数row_number())。


当前回答

你试过了吗

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)

删除MultiIndex并平化结果:

    id  value
0   1      1
1   1      2
2   2      1
3   2      2
4   3      1
5   4      1

其他回答

有时,提前排序整个数据非常耗时。 我们可以先分组,然后对每一组做topk:

g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)

你试过了吗

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)

删除MultiIndex并平化结果:

    id  value
0   1      1
1   1      2
2   2      1
3   2      2
4   3      1
5   4      1

要获得每个组的前N行,另一种方法是通过groupby().nth[:N]。此调用的结果与groupby().head(N)相同。例如,对于每个id的前2行,调用:

N = 2
df.groupby('id', as_index=False).nth[:N]

要获得每个组中最大的N个值,首先按“id”和“value”排序(确保通过适当地使用升序参数将id按升序排序,将“value”按降序排序),然后调用groupby().nth[]。

N = 2
df.sort_values(by=['id', 'value'], ascending=[True, False]).groupby('id', as_index=False).nth[:N]

这比在这里(1,2)的其他答案中建议的groupby().apply()调用快得多。在一个包含100k行和8000个组的示例中,%timeit测试表明它比这些解决方案快24-70倍。


此外,除了切片,你还可以将list/tuple/range传递给.nth()调用:

df.groupby('id', as_index=False).nth([0,1])
df.groupby('id', as_index=False).nth([0,2])  # <-- doesn't even have to be consecutive; 
                                             #     this returns 1st and 3rd row of each id

这适用于重复的值

如果你在前n个值中有重复的值,并且只想要唯一的值,你可以这样做:

import pandas as pd

ifile = "https://raw.githubusercontent.com/bhishanpdl/Shared/master/data/twitter_employee.tsv"
df = pd.read_csv(ifile,delimiter='\t')
print(df.query("department == 'Audit'")[['id','first_name','last_name','department','salary']])

    id first_name last_name department  salary
24  12   Shandler      Bing      Audit  110000
25  14      Jason       Tom      Audit  100000
26  16     Celine    Anston      Audit  100000
27  15    Michale   Jackson      Audit   70000

If we do not remove duplicates, for the audit department we get top 3 salaries as 110k,100k and 100k.
If we want to have not-duplicated salaries per each department, we can do this:

(df.groupby('department')['salary']
 .apply(lambda ser: ser.drop_duplicates().nlargest(3))
 .droplevel(level=1)
 .sort_index()
 .reset_index()
)

This gives

department  salary
0   Audit   110000
1   Audit   100000
2   Audit   70000
3   Management  250000
4   Management  200000
5   Management  150000
6   Sales   220000
7   Sales   200000
8   Sales   150000





从0.14.1开始,你现在可以在groupby对象上执行nmaximum和nleast:

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)