1

I am trying to select individual rows from a multiindex dataframe using a list of multiindices.

For example. I have got the following dataframe:

           Col1
A B C
1 1 1 -0.148593
    2  2.043589
  2 3 -1.696572
    4 -0.249049
2 1 5  2.012294
    6 -1.756410
  2 7  0.476035
    8 -0.531612

I would like to select all 'C' with (A,B) = [(1,1), (2,2)]

           Col1
A B C
1 1 1 -0.148593
    2  2.043589
2 2 7  0.476035
    8 -0.531612

My flawed code for this is as follows:

import pandas as pd
import numpy as np

arrays = [np.array([1, 1, 1, 1, 2, 2, 2, 2]), np.array([1, 1, 2, 2, 1, 1, 2, 2]), np.array([1, 2, 3, 4, 5, 6, 7, 8])]
df = pd.DataFrame(np.random.randn(8), index=arrays, columns=['Col1'])
df.rename_axis(['A','B','C'], inplace=True)
print(df)

idx_lst = [(1,1), (2,2)]
test = df.loc(axis=0)[idx_lst]
print(test)
jpp
  • 159,742
  • 34
  • 281
  • 339
Christian
  • 991
  • 2
  • 13
  • 25

2 Answers2

0

One option is to use pd.DataFrame.query:

res = df.query('((A == 1) & (B == 1)) | ((A == 2) & (B == 2))')

print(res)

           Col1
A B C          
1 1 1  0.981483
    2  0.851543
2 2 7 -0.522760
    8 -0.332099

For a more generic solution, you can use f-strings (Python 3.6+), which should perform better than str.format or manual concatenation.

filters = [(1,1), (2,2)]
filterstr = '|'.join(f'(A=={i})&(B=={j})' for i, j in filters)
res = df.query(filterstr)

print(filterstr)

(A==1)&(B==1)|(A==2)&(B==2)
jpp
  • 159,742
  • 34
  • 281
  • 339
0

The following might help:

idx_lst = [(1,1), (2,2)]

df.loc(0)[[ z for z in df.index if (z[0], z[1]) in idx_lst ]]
# Out[941]: 
#            Col1
# A B C          
# 1 1 1  0.293952
#     2  0.197045
# 2 2 7  2.007493
#     8  0.937420
jxc
  • 13,553
  • 4
  • 16
  • 34