Search code examples
pythonpandasmulti-index

Accessing a Series Based on Partial MultiIndex


I have a (very long) pandas Series with a 4-level MultiIndex, something like:

objs = ['car', 'bicycle', 'plane']
trials = ['trial A', 'trial B']
moves = [f'mov{i}' for i in range(1,11)]
multi = pd.MultiIndex.from_product([objs, trials, moves, range(10)], names=['obj', 'try', 'mov', 'time'])
s = pd.Series(np.random.rand(len(multi)))

I also have a list of tuples matching some of the indices under 'obj' and 'mov'. For example (in reality it's much longer):

idxs = [('car','mov1'), ('car','mov6'), ('plane','mov1')]

Is there a way to get all rows of s where the indices of multi contain those in idxs? I'm hoping to do this without re-indexing/slicing the Series because as I said it is very long and that would be inefficient.
I read the documentation for loc & xs, neither seem to be what I need.


Solution

  • One way to do it is to use list comprehension with pd.DataFrame.query:

    pd.concat([s.to_frame().query('obj == @i and mov == @j') for i, j in idxs])
    
    OR (as @Jon Nir Suggests which is faster)
    
    pd.concat([s.xs([i,j], level=[obj, mov]) for i,j in idxs])