基于2级列表对pandas 3级多索引进行切片

Question

这里有一个简单的例子：

import pandas as pd
import numpy as np
np.random.seed(0)
idx = pd.MultiIndex.from_product([[1,2,3], ['a', 'b', 'c'], [6, 7]])
df = pd.DataFrame(np.random.randn(18), index=idx)
selection = [(1, 'a'), (2, 'b')]

我想从 df 中选择所有索引以 selection 中的任意项开头的行。所以我想得到一个包含这些索引的 df 的子数据框：

(1, 'a', 6), (1, 'a', 7), (2, 'b', 6), (2, 'b', 7)

有什么最简单、最符合 Python 风格或 Pandas 风格的方法来做到这一点吗？我找到的内容是：

sel = [id[:2] in selection for id in df.index]
df.loc[sel]

Answer 1

试试这个：

print(pd.concat([df.loc[(*t, slice(None)), :] for t in selection]))

输出结果是：

              0
1 a 6  0.701403
    7 -0.822042
2 b 6 -0.610558
    7 -0.433880

Answer 2

一种选择是使用pyjanitor的 select 函数：

# pip install pyjanitor
import janitor
df.select(rows=selection)
              0
1 a 6 -0.989212
    7 -1.951114
2 b 6  1.629402
    7  0.298262

之所以这样做是因为元组的前两行和索引的前两层是相对应的。@mozways的解决方案更加通用。

Answer 3

你可以使用布尔索引，结合isin来进行选择：

out = df[df.index.isin(selection)]

输出结果：

              0
1 a 6  1.560268
    7  0.674709
2 b 6  0.848069
    7  0.130719

如果你想选择其他层级，可以去掉不需要的前面层级：

# here we want to select on levels 1 and 2
selection = [('a', 6), ('b', 7)]

df[df.index.droplevel(0).isin(selection)]

输出结果：

              0
1 a 6  1.560268
  b 7  0.137769
2 a 6  0.754946
  b 7  0.130719
3 a 6 -2.275646
  b 7 -2.199944