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

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

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 ())。

251217 次浏览

你试过了吗

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

自从0.14.1,你现在可以在groupby对象上执行nlargestnsmallest:

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上这样做,但目前它只适用于SeriesSeriesGroupBy。)

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

g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)
df.groupby('id').apply(lambda x : x.sort_values(by = 'value', ascending = False).head(2).reset_index(drop = True))
  • 这里的排序值升序为false,类似于nbiggest, True类似于nminimal。
  • 头部内的值与我们在nbiggest内给出的值相同,以获得每个组要显示的值的数量。
  • Reset_index是可选的,不是必需的。

这适用于重复的值

如果你在前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










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

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

要得到N的最大值,首先按"id"和“;value"(请确保将id按升序排序,“;value"依次降序使用ascending参数),然后调用groupby().nth[]

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

这比这里的其他答案(12)中建议的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