python - 如何使用 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

现在我想获取如下索引:

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

索引 i4 在 [1, 2] 中索引 i1 = 2

                 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 中使用列表理解来获得你的最终目标。下面，假设 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])