Эффективное нахождение строк, следующих за подмножеством строк DataFrame панелей MultiIndex - PullRequest
0 голосов
/ 04 сентября 2018

Как можно эффективно найти (т. Е. Векторизованное решение) строки, которые следуют за подмножеством строк в палитре данных 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

Спасибо за вашу помощь!

1 Ответ

0 голосов
/ 05 сентября 2018

Если вы хотите выбрать следующие строки некоторого произвольного подмножества, вы можете сделать это, создав маску:

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