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

Question

我有一个4级MultiIndex的pandas数据帧。 我正在尝试为每个级别1索引选择具有不同级别4索引的行。

例：

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

现在我想获得索引如下：

对于索引i1 = 1，索引i4在[0,1]中

对于索引i1 = 2，[1,2]中的索引i4

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

现在这个工作：

    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]

但它看起来不好解决方案。 这有什么聪明的方法吗？

Answer 1

使用IndexSlice可以使这更容易，如下所示：

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

如果您需要创建许多索引，则可以将这些索引存储在数据框中并迭代该数据框以创建各种切片，然后使用pd.concat的列表pd.concat来获取最终对象。 下面，假设x ['id1']是你想要id1拥有的值，我还假设你要限制相同的两个索引列。

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])