通过multiindex的子集从pandas中选择行

时间:2014-03-20 01:00:36

标签: python pandas

我在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而失败。

This question似乎相似,但由于index.get_level_values仅适用于单个级别而非多个级别,因此其中的答案不适用于除顶级索引之外的任何其他内容。

根据建议here我设法完成了我想要的事情

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

然而,使用len(set(index.get_level_values(...)))获取不同值的计数并通过迭代遍历每一行来构建布尔向量感觉更像是我在框架中实现看似简单的任务多索引设置。有更好的解决方案吗?

这是使用pandas 0.13.1。

1 个答案:

答案 0 :(得分:2)

可能会有更好的内容,但您至少可以通过使用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]

它比我之前的建议更快。它确实适用于小型DataFrame:

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

并且对于大型DataFrame来说仍然是最快的,但不是那么多:

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