Pandas: выберите первую пару строк в каждой группе

Question

Я не могу решить эту простую проблему, и я прошу помочь здесь ... У меня есть DataFrame следующим образом, и я хочу выбрать первые две строки в каждой группа «в»

df = pd.DataFrame({'a':pd.Series(['NewYork','NewYork','NewYork','Washington','Washington','Texas','Texas','Texas','Texas']), 'b': np.arange(9)}) 

df 
Out[152]: 
      a b 
0  NewYork 0 
1  NewYork 1 
2  NewYork 2 
3 Washington 3 
4 Washington 4 
5  Texas 5 
6  Texas 6 
7  Texas 7 
8  Texas 8 

, то есть, я хочу выход следующим образом:

  a b 
0  NewYork 0 
1  NewYork 1 
2 Washington 3 
3 Washington 4 
4  Texas 5 
5  Texas 6 

Спасибо большое за помощь.

Answer 1

К сожалению, это, кажется, кто-то задал подобный вопрос перед ... Pandas dataframe get first row of each group Понял теперь ...

df.groupby('a').head(2).reset_index(drop=True) 
Out[165]: 
      a b 
0  NewYork 0 
1  NewYork 1 
2  Texas 5 
3  Texas 6 
4 Washington 3 
5 Washington 4 

Answer 2

В панд 0.13rc, вы можете сделать это непосредственно с помощью головки (т.е. нет необходимости к reset_index):

In [11]: df.groupby('id', as_index=False).head(2) 
Out[11]: 
    id value 
0 1 first 
1 1 second 
3 2 first 
4 2 second 
5 3 first 
6 3 third 
9 4 second 
10 4 fifth 
11 5 first 
12 6 first 
13 6 second 
15 7 fourth 
16 7 fifth 

[13 rows x 2 columns] 

Примечания: правильные показателям, и это значительно быстрее, чем раньше (с или без reset_index), даже с этим небольшим примером:

# 0.13rc 
In [21]: %timeit df.groupby('id', as_index=False).head(2) 
1000 loops, best of 3: 279 µs per loop 

# 0.12 
In [21]: %timeit df.groupby('id', as_index=False).head(2) # this didn't work correctly 
1000 loops, best of 3: 1.76 ms per loop 

In [22]: %timeit df.groupby('id').head(2).reset_index(drop=True) 
1000 loops, best of 3: 1.82 ms per loop