如何使用MultiIndex的相关级别对MultiIndex DataFrame进行切片

Question

I have a pandas dataframe with 4 levels of a MultiIndex. I am trying to select rows which has different level 4 indexes for each level 1 indexses.

example:

In [68]: df = pd.DataFrame({'i1':[1,1,1,2,2,2],
                        'i2':[1,1,2,1,1,2],
                        'i3':[1,1,1,1,1,1],
                        'i4':[0,1,2,0,1,2],
                        'data':[1,1,2,2,1,1]}).set_index(['i1','i2','i3','i4'])


In [69]: df
Out[69]:
             data
i1 i2 i3 i4
1  1  1  0      1
         1      1
   2  1  2      2
2  1  1  0      2
         1      1
   2  1  2      1

Now I want to get indexses as follows:

index i4 in [0, 1] for index i1 = 1

index i4 in [1, 2] for index i1 = 2

                 data
i1 i2 i3 i4
1  1  1  0      1
         1      1
2  1  1  1      1
   2  1  2      1

For now this works:

    cond1 = (df.index.get_level_values('i1') == 1) & (df.index.get_level_values('i4').isin([0,1]))
    cond2 = (df.index.get_level_values('i1') == 2) & (df.index.get_level_values('i4').isin([1,2]))
    .
    .
    .
    condN = ...
    df[cond1 | cond2 | ... | condN]

but it looks like bad solution. Is there any clever way of doing this?

Answer 1

You can make this a bit easier with IndexSlice , as follows:

idx = pd.IndexSlice
index1 = idx[1, :, :, 0:1]
index2 = idx[2, :, :, 1:2]
pd.concat([df.loc[index1], df.loc[index2]])

If you have many indices you need to create, you can store those indices in a dataframe and iterate over that dataframe to create your various slices and then use a list comprehension in pd.concat to get your final object. Below, x['id1'] is assumed to be the value you want id1 to have, and I also make the assumption that you want to limit the same two index columns.

indices = [
    idx[
        x['id1'],
        lambda x['id2']: x['id2'] or slice(None),
        lambda x['id3']: x['id3'] or slice(None),
        x['id4']
    ] for x in index_df.iterrows()
]
pd.concat([df.loc[i] for i in indices])