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

wfauudbj  于 2023-02-14  发布在  其他
关注(0)|答案(6)|浏览(134)

假设我有这样的PandasDataFrame:

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,其中每个ID有前2条记录,如下所示:

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

2exbekwf

2exbekwf1#

你有没有试

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

(Keep请记住,根据您的数据,您可能需要在此之前进行排序/排序)
编辑:正如提问者提到的,使用

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
stszievb

stszievb2#

Since 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)来完全摆脱它。
(Note:From 0.17.1您也可以在DataFrameGroupBy上执行此操作,但目前它仅适用于SeriesSeriesGroupBy。)

cmssoen2

cmssoen23#

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

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

oug3syen4#

df.groupby('id').apply(lambda x : x.sort_values(by = 'value', ascending = False).head(2).reset_index(drop = True))
  • 在这里,按false升序对值进行排序,得到类似于nlargest的值,按True升序对值排序,得到类似于nminest的值。
  • head中的值与我们在nlargest中给予的值相同,nlargest用于获取每个组显示的值的数量。
  • reset_index是可选的而不是必需的。
lzfw57am

lzfw57am5#

这适用于重复值

如果前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
x6yk4ghg

x6yk4ghg6#

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

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

为了获得每组的最大N值,我建议两种方法。
1.首先按“id”和“value”排序(通过适当使用ascending参数,确保“id”按升序排序,“value”按降序排序),然后调用groupby().nth[]

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

1.另一种方法是对每个组的值进行排序,并使用这些排序进行过滤。

# for the entire rows
N = 2
msk = df.groupby('id')['value'].rank(method='first', ascending=False) <= N
df1 = df[msk]

# for specific column rows
df1 = df.loc[msk, 'value']

这两种方法都比groupby().apply()groupby().nlargest()调用快得多,正如这里的其他答案(123)所建议的那样。在一个具有100 k行和8000个组的样本上,%timeit测试显示它比这些解决方案快24-150倍。
此外,除了切片,您还可以将列表/元组/范围传递给.nth()调用:

df.groupby('id', as_index=False).nth([0,1])

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

相关问题