熊猫在每组中获得最高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  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

我想为每个id获取一个包含前2条记录的新DataFrame,如下所示:

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

我可以通过以下方式在组内编号记录:

>>> dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()
>>> dfN
   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()).

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.)

Answer 3

有时，提前对整个数据进行排序非常耗时。我们可以先分组，然后为每个组做topk：

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

Answer 4

df.groupby('id').apply(lambda x : x.sort_values(by = 'value', ascending = False).head(2).reset_index(drop = True))

• 这里排序值升序 false 给出类似于 nlargest 的结果，True 给出类似于 nsmallest 的结果。
• head 内的值与我们在 nlargest 内给出的值相同，以获得每个组要显示的值的数量。
• reset_index 是可选的，但不是必需的。

Answer 5

要获取每组的前 N ​​行，另一种方法是 via 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]
   

   
   

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

   # 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()调用快得多（1,2,3 ）。在具有 100k 行和 8000 组的样本上，测试表明它比那些解决方案快 24-150 倍。groupby().nlargest()%timeit

此外，您还可以将列表/元组/范围传递给调用，而不是切片.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])

Answer 6

这适用于重复的值

如果您在 top-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