有效地查找MultiIndex pandas DataFrame的行子集之后的行

时间:2018-09-04 13:14:14

标签: python pandas multi-index shift

如何有效地找到(即矢量化解决方案)紧随MultiIndex pandas DataFrame的行子集的行?

对于单个索引,似乎可以使用pandas.Index.shift

示例:

import pandas as pd

# original data-frame
t = pd.DataFrame(data={'i1':[0,0,0,0,1,1,1,1,2,2,2,2],
                       'i2':[0,1,2,3,0,1,2,3,0,1,2,3],
                       'x':[1.,2.,3.,4.,5.,6.,7.,8.,9.,10.,11.,12.]})
t.set_index(['i1','i2'], inplace=True)
t.sort_index(inplace=True)
print(t)

# subset of rows
t2 = t.loc[(slice(None),slice(1,1)),:]
print(t2)

# example of *not efficient* solution (i.e. not vectorized)
t3 = t.iloc[ [t.index.get_loc(v)+1 for v in t2.index] ]
print(t3)
# original DataFrame
          x
i1 i2      
0  0    1.0
   1    2.0
   2    3.0
   3    4.0
1  0    5.0
   1    6.0
   2    7.0
   3    8.0
2  0    9.0
   1   10.0
   2   11.0
   3   12.0

# subset of rows
          x
i1 i2      
0  1    2.0
1  1    6.0
2  1   10.0

# expected solution
          x
i1 i2      
0  2    3.0
1  2    7.0
2  2   11.0

谢谢您的帮助!

1 个答案:

答案 0 :(得分:1)

如果要选择任意子集的以下行,可以通过创建遮罩来实现:

mask = pd.Series(False, index=t.index)
mask[t2.index] = True

然后,您可以使用移位的掩码为t编制索引:

t3 = t.loc[mask.shift(1).fillna(False)]
# and maybe:
t4 = t.loc[mask.shift(2).fillna(False)]

但是,这听起来像是XY问题。你真正想要的是什么?如果您只想方便地在第二级多索引上建立索引,则应尝试IndexSlice

idx = pd.IndexSlice
t2 = t.loc[idx[:,1],:]
t3 = t.loc[idx[:,2],:]