Выберите из DataFrame на основе нескольких уровней MultiIndex - PullRequest
Новая
       8

Выберите из DataFrame на основе нескольких уровней MultiIndex

0 голосов
/ 05 декабря 2018

Как расширить логику выбора из DataFrame на основе первых N-1 уровней, когда N> 2?

В качестве примера рассмотрим DataFrame: 

midx = pd.MultiIndex.from_product([[0, 1], [10, 20, 30], ["a", "b"]])
df = pd.DataFrame(1, columns=midx, index=np.arange(3))
In[11]: df
Out[11]: 
   0                 1               
  10    20    30    10    20    30   
   a  b  a  b  a  b  a  b  a  b  a  b
0  1  1  1  1  1  1  1  1  1  1  1  1
1  1  1  1  1  1  1  1  1  1  1  1  1
2  1  1  1  1  1  1  1  1  1  1  1  1

Здесь, легко выбрать столбцы, где 0 или 1 находятся на первом уровне: 

df[[0, 1]]

Но та же логика не распространяется на выбор столбцов с 0 или 1 в первом и 10 или 20 во второмуровень: 

In[13]: df[[(0, 10), (0, 20), (1, 10), (1, 20)]]
ValueError: operands could not be broadcast together with shapes (4,2) (3,) (4,2)

Работает: 

df.loc[:, pd.IndexSlice[[0, 1], [10, 20], :]]

, но это громоздко, особенно когда селектор нужно извлечь из другого DataFrame с двухуровневым MultiIndex: 

idx = df.columns.droplevel(2)
In[16]: idx
Out[16]: 
MultiIndex(levels=[[0, 1], [10, 20, 30]],
           labels=[[0, 0, 0, 0, 0, 0, 1, 1, 1, ... 1, 2, 2]])
In[17]: df[idx]
ValueError: operands could not be broadcast together with shapes (12,2) (3,) (12,2)

РЕДАКТИРОВАТЬ: В идеале я также хотел бы иметь возможность упорядочивать столбцы таким образом, а не просто выбирать их - опять же, в духе df[[1, 0]] возможность упорядочить столбцы на основепервый уровень.

1 Ответ

0 голосов
/ 05 декабря 2018

Если возможно, вы можете отфильтровать по boolean indexing с get_level_values и isin: 

m1 = df.columns.get_level_values(0).isin([0,1])
m2 = df.columns.get_level_values(1).isin([10,20])

print (m1)
[ True  True  True  True  True  True  True  True  True  True  True  True]
print (m2)
[ True  True  True  True False False  True  True  True  True False False]
print (m1 & m2)
[ True  True  True  True False False  True  True  True  True False False]

df1 = df.loc[:, m1 & m2]
print (df1)
   0           1         
  10    20    10    20   
   a  b  a  b  a  b  a  b
0  1  1  1  1  1  1  1  1
1  1  1  1  1  1  1  1  1
2  1  1  1  1  1  1  1  1

df.columns = df.columns.droplevel(2)
print (df)
   0                 1               
  10 10 20 20 30 30 10 10 20 20 30 30
0  1  1  1  1  1  1  1  1  1  1  1  1
1  1  1  1  1  1  1  1  1  1  1  1  1
2  1  1  1  1  1  1  1  1  1  1  1  1

df2 = df.loc[:, m1 & m2]
print (df2)
   0           1         
  10 10 20 20 10 10 20 20
0  1  1  1  1  1  1  1  1
1  1  1  1  1  1  1  1  1
2  1  1  1  1  1  1  1  1
...