python - 如何过滤前 N 项的 groupby

Question

在 Pandas 中，如何修改 groupby 以仅获取组中的前 N ​​个项目？

示例

df = pd.DataFrame({'id': [1, 1, 1, 2, 2, 2, 2], 
                   'values': [1, 2, 3, 4, 5, 6, 7]})
>>> df
   id  values
0   1       1
1   1       2
2   1       3
3   2       4
4   2       5
5   2       6
6   2       7

所需功能

# This doesn't work, but I am trying to return the first two items per group.
>>> df.groupby('id').first(2)  
   id  values
0   1       1
1   1       2
3   2       4
4   2       5

我尝试过的

我可以执行 groupby 并迭代组以获取前 n 个值的索引，但必须有一个更简单的解决方案。

n = 2  # First two rows.
idx = [i for group in df.groupby('id').groups.itervalues() for i in group[:n]]
>>> df.ix[idx]
   id  values
0   1       1
1   1       2
3   2       4
4   2       5

Answer 1

您可以使用head :

In [11]: df.groupby("id").head(2)
Out[11]:
   id  values
0   1       1
1   1       2
3   2       4
4   2       5

注意:在旧版本中，这曾经相当于 .apply(pd.DataFrame.head) 但从 0.15 开始它的效率更高(？)，现在它使用 cumcount在引擎盖下。