Как можно эффективно найти (т. Е. Векторизованное решение) строки, которые следуют за подмножеством строк в палитре данных MultiIndex pandas?
Кажется, что для одного индекса можно использовать 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
Спасибо за вашу помощь!