Select rows in pandas MultiIndex DataFrame by giving values for several (but not all) levels

Viewed 78
m_index = pd.MultiIndex.from_product([['a', 'b', 'c'], [1, 2, 3], [0.1, 0.2, 0.3]])
test_df = pd.DataFrame(np.arange(27), m_index)

I want to find all the rows for which the first two index levels are in some list of values, e.g. [('a', 1), ('b', 2), ('c', 1), ('c', 3)]

Is this possible?

3 Answers

You can use droplevel on the index and then call isin:

keys = [('a', 1), ('b', 2), ('c', 1), ('c', 3)]
# -1 drops the last level, so we're checking the first two only
test_df[test_df.index.droplevel(-1).isin(keys)]

          0
a 1 0.1   0
    0.2   1
    0.3   2
b 2 0.1  12
    0.2  13
    0.3  14
c 1 0.1  18
    0.2  19
    0.3  20
  3 0.1  24
    0.2  25
    0.3  26

Use DataFrame.reset_index with Index.isin and filter in boolean indexing:

L = [('a', 1), ('b', 2), ('c', 1), ('c', 3)]

df = test_df[test_df.reset_index(level=2).index.isin(L)]
print (df)
          0
a 1 0.1   0
    0.2   1
    0.3   2
b 2 0.1  12
    0.2  13
    0.3  14
c 1 0.1  18
    0.2  19
    0.3  20
  3 0.1  24
    0.2  25
    0.3  26
    

A solution is also possible without index manipulation. Just match the zipped tuple of index level 0 and 1 to the desired list.

test_df[list(map(lambda el: el in [('a', 1), ('b', 2), ('c', 1), ('c', 3)],
                 zip(test_df.index.get_level_values(0),
                     test_df.index.get_level_values(1))))]

          0
a 1 0.1   0
    0.2   1
    0.3   2
b 2 0.1  12
    0.2  13
    0.3  14
c 1 0.1  18
    0.2  19
    0.3  20
  3 0.1  24
    0.2  25
    0.3  26
Related