高效地查找多索引熊猫数据帧的行子集后面的行

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