Группировка и фильтрация на основе МАКСИМАЛЬНОГО значения одного столбца И ОБЯЗАТЕЛЬНО topN значений другого столбца - PullRequest
Новая
       4

Группировка и фильтрация на основе МАКСИМАЛЬНОГО значения одного столбца И ОБЯЗАТЕЛЬНО topN значений другого столбца

1 голос
/ 02 февраля 2020

У меня есть следующие данные: 

    individual  groupID  choice     probA     probB
0      9710535        0       0  0.025589  0.008647
1      9710535        2       0  0.036252  0.014684
2      9710535        2       0  0.036252  0.013894
3      9710535        1       0  0.090057  0.030648
4      9710535        1       0  0.090057  0.014804
5      9710535        5       0  0.512675  0.021444
6      9710535        5       0  0.512675  0.020298
7      9710535        5       0  0.512675  0.163837
8      9710535        5       0  0.512675  0.085191
9      9710556        1       1  0.034381  0.796679
10     9710556        1       0  0.034381  0.796679
11     9710556        1       0  0.034381  0.796679
12     9710556        1       0  0.034381  0.157622
13     9710556        1       0  0.034381  0.157622
14     9710556        0       0  0.025589  0.008647

Я хочу отфильтровать данные на основе максимального значения столбца для каждого individual на probA для каждого groupID, что составляет groupID = 5 для individual 9710535. Однако, если размер этой группы меньше 5, как в данном случае (в ней всего 4 элемента), я также хочу, чтобы строки из других групп были верхними (5 - размер этой группы) на основе столбца probB .

Конечный результат должен выглядеть следующим образом: 

  individual  groupID  choice     probA     probB
0     9710535        1       0  0.090057  0.030648
1     9710535        5       0  0.512675  0.021444
2     9710535        5       0  0.512675  0.020298
3     9710535        5       0  0.512675  0.163837
4     9710535        5       0  0.512675  0.085191
5     9710556        1       1  0.034381  0.796679
6     9710556        1       0  0.034381  0.796679
7     9710556        1       0  0.034381  0.796679
8     9710556        1       0  0.034381  0.157622
9     9710556        1       0  0.034381  0.157622

Я работаю над подходом к пользовательской функции с этим псевдокодом: 

def custom_filter(df, groub_by = 'individual', 
                  condition_column1 = 'probA', 
                  condition_column2 = 'probB',
                  top_count = 5)
    return filtered_df

Буду признателен за любую помощь !

1 Ответ

2 голосов
/ 02 февраля 2020

Я думаю, вам нужно сначала извлечь все максимальные строки на GroupBy.transform с max и сравнить groupID с фильтром по boolean indexing: 

mask = df['groupID'].eq(df.groupby('individual')['groupID'].transform('max'))
df1 = df[mask]
print (df1)
    individual  groupID  choice     probA     probB
5      9710535        5       0  0.512675  0.021444
6      9710535        5       0  0.512675  0.020298
7      9710535        5       0  0.512675  0.163837
8      9710535        5       0  0.512675  0.085191
9      9710556        1       1  0.034381  0.796679
10     9710556        1       0  0.034381  0.796679
11     9710556        1       0  0.034381  0.796679
12     9710556        1       0  0.034381  0.157622
13     9710556        1       0  0.034381  0.157622

Затем отфильтруйте несопоставленные строки по Index.isin и отсортируйте по 2 столбцам с помощью DataFrame.sort_values: 

df2 = (df[~df.index.isin(df1.index)]
           .sort_values(['individual','probB'], 
                        ascending=[True, False]))
print (df2)
    individual  groupID  choice     probA     probB
3      9710535        1       0  0.090057  0.030648
4      9710535        1       0  0.090057  0.014804
1      9710535        2       0  0.036252  0.014684
2      9710535        2       0  0.036252  0.013894
0      9710535        0       0  0.025589  0.008647
14     9710556        0       0  0.025589  0.008647

Последнее объединение по concat и получить верхнюю 5 по GroupBy.head и последнюю сортировку по другим 2 столбцам: 

df = (pd.concat([df1, df2])
        .groupby('individual').head()
        .sort_values(['individual','groupID']))
print (df)
    individual  groupID  choice     probA     probB
3      9710535        1       0  0.090057  0.030648
5      9710535        5       0  0.512675  0.021444
6      9710535        5       0  0.512675  0.020298
7      9710535        5       0  0.512675  0.163837
8      9710535        5       0  0.512675  0.085191
9      9710556        1       1  0.034381  0.796679
10     9710556        1       0  0.034381  0.796679
11     9710556        1       0  0.034381  0.796679
12     9710556        1       0  0.034381  0.157622
13     9710556        1       0  0.034381  0.157622
...