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

Question

如何有效地找到跟随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

谢谢您的帮助！

Answer 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],:]