通过多层索引子集选择pandas行

Question

我在pandas中有一个多重索引的数据框，索引里有4列，还有一些数据列。下面是一个例子：

import pandas as pd
import numpy as np
cnames = ['K1', 'K2', 'K3', 'K4', 'D1', 'D2']
rdata = pd.DataFrame(np.random.randint(1, 3, size=(8, len(cnames))), columns=cnames)
rdata.set_index(cnames[:4], inplace=True)
rdata.sortlevel(inplace=True)
print(rdata)
             D1  D2
K1 K2 K3 K4        
1  1  1  1    1   2
         1    1   2
      2  1    2   1
   2  1  2    2   1
      2  1    2   1
2  1  2  2    2   1
   2  1  2    1   1
         2    1   1

[8 rows x 2 columns]

我想做的是选择那些在K3层级上恰好有2个不同值的行。不是2行，而是两个不同的值。我找到了一种方法来生成我想要的那种“掩码”：

filterFunc = lambda x: len(set(x.index.get_level_values('K3'))) == 2
mask = rdata.groupby(level=cnames[:2]).apply(filterFunc)
print(mask)
K1  K2
1   1      True
    2      True
2   1     False
    2     False
dtype: bool

我本以为既然 rdata.loc[1, 2] 可以只匹配索引的一部分，那么用这样的布尔向量也应该可以做到。不幸的是， rdata.loc[mask] 失败了，出现了 IndexingError: Unalignable boolean Series key provided 的错误。

这个问题看起来很相似，但那里的答案只适用于顶层索引，因为 index.get_level_values 只能在单一层级上工作，而不能在多个层级上使用。

根据 这里的建议，我成功地用

rdata[[mask.loc[k1, k2] for k1, k2, k3, k4 in rdata.index]]

实现了我想要的功能，不过，使用 len(set(index.get_level_values(...))) 来计算不同值的数量，以及之后通过遍历每一行来构建布尔向量，感觉就像是在和框架作斗争，而这在多重索引的设置中似乎应该是个简单的任务。有没有更好的解决方案呢？

这是在使用pandas 0.13.1。

Answer 1

可能还有更好的方法，但你至少可以通过使用groupby-filter来省去定义mask这一步：

rdata.groupby(level=cnames[:2]).filter(
      lambda grp: (grp.index.get_level_values('K3')
                      .unique().size) == 2)

Out[83]: 
             D1  D2
K1 K2 K3 K4        
1  1  1  1    1   2
         1    1   2
      2  1    2   1
   2  1  2    2   1
      2  1    2   1

[5 rows x 2 columns]

这个方法比我之前提到的建议要快。对于小型数据框，它表现得非常好：

In [84]: %timeit rdata.groupby(level=cnames[:2]).filter(lambda grp: grp.index.get_level_values('K3').unique().size == 2)
100 loops, best of 3: 3.84 ms per loop

In [76]: %timeit rdata2.groupby(level=cnames[:2]).filter(lambda grp: grp.groupby(level=['K3']).ngroups == 2)
100 loops, best of 3: 11.9 ms per loop

In [77]: %timeit rdata2.groupby(level=cnames[:2]).filter(lambda grp: len(set(grp.index.get_level_values('K3'))) == 2)
100 loops, best of 3: 13.4 ms per loop

而对于大型数据框，它仍然是最快的，虽然速度差距没有那么大：

In [78]: rdata2 = pd.concat([rdata]*100000)

In [85]: %timeit rdata2.groupby(level=cnames[:2]).filter(lambda grp: grp.index.get_level_values('K3').unique().size == 2)
1 loops, best of 3: 756 ms per loop

In [79]: %timeit rdata2.groupby(level=cnames[:2]).filter(lambda grp: grp.groupby(level=['K3']).ngroups == 2)
1 loops, best of 3: 772 ms per loop

In [80]: %timeit rdata2.groupby(level=cnames[:2]).filter(lambda grp: len(set(grp.index.get_level_values('K3'))) == 2)
1 loops, best of 3: 1 s per loop