从前两组获得前两行

Question

Consider the data frame df

mux = pd.MultiIndex.from_arrays([
    list('aaaabbbbbccdddddd'),
    list('tuvwlmnopxyfghijk')
], names=['one', 'two'])

df = pd.DataFrame({'col': np.arange(len(mux))}, mux)

df

         col
one two     
a   t      0
    u      1
    v      2
    w      3
b   l      4
    m      5
    n      6
    o      7
    p      8
c   x      9
    y     10
d   f     11
    g     12
    h     13
    i     14
    j     15
    k     16

How do I elegantly get the first two rows of the first two groups if I group by the first level of the index:

         col
one two     
a   t      0
    u      1
b   l      4
    m      5

Answer 1

Option 1
You could use a list comp and pd.concat :

pd.concat([g.head(2) for _, g in df.groupby(level=0)][:2])

         col
one two     
a   t      0
    u      1
b   l      4
    m      5

Since having the list comp complete is an unnecessary overhead, you could use itertools.takewhile to prevent that.

it = itertools.takewhile(lambda x: x[0] < 2, enumerate(df.groupby(level=0)))
pd.concat([g.head(2) for _, (_, g) in it])

         col
one two     
a   t      0
    u      1
b   l      4
    m      5

---

Option 2
Another possible solution I could think of is pre-filtering your df to retain rows for only the first two values of index level 0, and then do the groupby.

# https://stackoverflow.com/a/46900625/4909087
df.loc[df.index.levels[0][:2].values].groupby(level=0).head(2)

         col
one two     
a   t      0
    u      1
b   l      4
    m      5

Answer 2

Looks hacky but this is what I tried

df.groupby(level=['one']).head(2)[:4]


       col
one two 
a   t   0
u   1
b   l   4
m   5