Pandas在每组中获得最多n项记录

aydmsdu9  于 2022-09-21  发布在  其他
关注(0)|答案(5)|浏览(177)

假设我有这样的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有前两条记录,如下所示:

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

iyr7buue

iyr7buue1#

你试过了吗

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)

要删除多重索引并展平结果,请执行以下操作:

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

mspsb9vt2#

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)来彻底摆脱它。

(注意:From 0.17.1您也可以在DataFrameGroupBy上执行此操作,但目前它仅适用于SeriesSeriesGroupBy。)

iyfjxgzm

iyfjxgzm3#

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

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

sdnqo3pr4#

df.groupby('id').apply(lambda x : x.sort_values(by = 'value', ascending = False).head(2).reset_index(drop = True))
  • 在这里,排序值按升序FALSE表示类似于n最大,按True升序表示类似于N最小。
  • Head中的值与我们在nmax中提供的值相同,以获取要为每个组显示的值的数量。
  • RESET_INDEX是可选的,不是必需的。
apeeds0o

apeeds0o5#

对重复值有效

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

相关问题