Условная группировка фрейма данных Pandas

Question

У меня есть фрейм данных Pandas вроде:

  ID  Group   Rank include 
  1   Group1  1    0     
  2   Group1  2    1     
  3   Group1  3    0     
  4   Group1  4    1     
  5   Group2  1    0     
  6   Group2  2    0     
  7   Group2  3    1    
  8   Group2  4    1     
  9   Group2  5    1      

Я хочу выбрать подмножество, которое удовлетворяет следующим условиям: - максимум 3 ряда из группы - включить все строки, которые включают include = 1 (их не более 3 на группу) - минимизировать общий рейтинг по группе

По сути, я хочу выбрать TOP "N" записей из каждой группы, но строки с include = 1 имеют приоритет над остальными.

Фрейм выходных данных:

  ID  Group   Rank include 
  1   Group1  1    0     
  2   Group1  2    1       
  4   Group1  4    1         
  7   Group2  3    1    
  8   Group2  4    1     
  9   Group2  5    1

Я пробовал несколько подходов, но ни один из них не сработал, и теперь я застрял.
Пожалуйста, помогите!

Answer 1

Сначала отсортируйте фрейм данных по вашим условиям:

df.sort_values(['include', 'Rank'], ascending=[False, True], inplace=True)

Затем сгруппируйте Group и получите три верхних строки для каждой группы. Кроме того, добавьте sort_values для сортировки окончательного фрейма данных по Group, а затем по Rank, например:

df.groupby('Group').head(3).sort_values(['Group', 'Rank'])

Выход:

   ID   Group  Rank  include
0   1  Group1     1        0
1   2  Group1     2        1
3   4  Group1     4        1
6   7  Group2     3        1
7   8  Group2     4        1
8   9  Group2     5        1