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在幕后。