Как выбрать строки из многоиндексного фрейма данных на основе условия в одном столбце

#python #pandas #dataframe

Вопрос:

У меня есть многоиндексный фрейм данных, который выглядит так:

 df = {'C': {('S', 0): 'A',
  ('S', 2): 'A',
  ('T', 0): 'A',
  ('T', 1): 'A',
  ('T', 3): 'A',
  ('U', 1): 'A',
  ('U', 2): 'A',
  ('U', 0): 'A',
  ('V', 0): 'A',
  ('W', 2): 'A',
  ('W', 0): 'A',
  ('X', 0): 'A',
  ('Y', 3): 'A',
  ('Z', 0): 'A',
  ('Z', 1): 'A'},
 'D': {('S', 0): '15',
  ('S', 2): '22',
  ('T', 0): '20',
  ('T', 1): '20',
  ('T', 3): '20',
  ('U', 1): '18',
  ('U', 2): '14',
  ('U', 0): '14',
  ('V', 0): '14',
  ('W', 2): '22',
  ('W', 0): '25',
  ('X', 0): '15',
  ('Y', 3): '17',
  ('Z', 0): '04',
  ('Z', 1): '16'},
 'E': {('S', 0): 1.0,
  ('S', 2): 1.0,
  ('T', 0): 2.0,
  ('T', 1): 2.0,
  ('T', 0): 2.0,
  ('U', 1): 2.0,
  ('U', 2): 2.0,
  ('U', 0): 2.0,
  ('V', 0): 1.0,
  ('W', 2): 1.0,
  ('W', 0): 1.0,
  ('X', 0): 1.0,
  ('Y', 3): 2.0,
  ('Z', 0): 3.0,
  ('Z', 1): 3.0}}

 

Я хочу сохранить строки 0-го уровня, если на его уровне 1 есть значение >=2

результат будет выглядеть примерно так:

 outp = {'C': {('S', 0): 'A',
  ('S', 2): 'A',
  ('T', 0): 'A',
  ('T', 1): 'A',
  ('T', 3): 'A',
  ('U', 1): 'A',
  ('U', 2): 'A',
  ('U', 0): 'A',
  ('W', 2): 'A',
  ('W', 0): 'A',
  ('Y', 3): 'A'},
 'D': {('S', 0): '15',
  ('S', 2): '22',
  ('T', 0): '20',
  ('T', 1): '20',
  ('T', 3): '20',
  ('U', 1): '18',
  ('U', 2): '14',
  ('U', 0): '14',
  ('W', 2): '22',
  ('W', 0): '25',
  ('Y', 3): '17'},
 'E': {('S', 0): 1.0,
  ('S', 2): 1.0,
  ('T', 0): 2.0,
  ('T', 1): 2.0,
  ('T', 0): 2.0,
  ('U', 1): 2.0,
  ('U', 2): 2.0,
  ('U', 0): 2.0,
  ('W', 2): 1.0,
  ('W', 0): 1.0,
  ('Y', 3): 2.0}}
 

Что я сделал, так это то, что я получил значение с уровня 0, когда уровень 1 >= 2, но поскольку при этом я удалил значения 0 и 1 с уровня 1, которые должны остаться, мне пришлось создать другой фрейм данных с полученными значениями, а затем объединить, используя «внутренний». Я получил желаемый результат, но наверняка выбрал долгий и, вероятно, глупый путь.

Как я мог бы сделать это лучше?

Спасибо.

Ответ №1:

Давайте попробуем groupby filter включить level=0 и отфильтровать, чтобы сохранить значения уровня 0 any , когда значение уровня индекса 1 ( get_level_values ) больше или равно 2:

 outp = (
    df.groupby(level=0)
        .filter(lambda s: (s.index.get_level_values(1) >= 2).any())
)
 

outp :

      C   D    E
S 0  A  15  1.0
  2  A  22  1.0
T 0  A  20  2.0
  1  A  20  2.0
  3  A  20  NaN
U 1  A  18  2.0
  2  A  14  2.0
  0  A  14  2.0
W 2  A  22  1.0
  0  A  25  1.0
Y 3  A  17  2.0
 

Ответ №2:

Получите индексы на уровне 0, где уровень 1 >= 2, и проиндексируйте основные > df :

 df.loc[df.query("ilevel_1 > =2").index.get_level_values(0)]

     C   D    E
S 0  A  15  1.0
  2  A  22  1.0
T 0  A  20  2.0
  1  A  20  2.0
  3  A  20  NaN
U 1  A  18  2.0
  2  A  14  2.0
  0  A  14  2.0
W 2  A  22  1.0
  0  A  25  1.0
Y 3  A  17  2.0
 

Ответ №3:

Вот другой способ:

 (df.loc[df.reset_index(level=1)
        .groupby(level=0)['level_1']
        .transform(lambda x: x.ge(2).any()).to_numpy()])