问题描述：

假设我有这样的 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]})

看起来像：

   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

我可以使用以下方式对组内的记录进行编号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()）。

解决方案 1：

你试过吗

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

解决方案 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。）

解决方案 3：

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

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

解决方案 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 是可选的，不是必需的。

解决方案 5：

要获取每个组的前 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']

这两个方法都比这里其他答案( 1、2、3groupby().apply() ) 中建议的调用快得多。在一个有 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])

解决方案 6：

这适用于重复的值

如果前 n 个值中有重复的值，并且只想要唯一值，则可以这样做：

import pandas as pd

ifile = "https://raw.githubusercontent.com/bhishanpdl/Shared/master/data/twitter_employee.tsv"
df = pd.read_csv(ifile,delimiter='    ')
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