Как фильтровать groupby для первых N элементов

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 под капотом (?).