Группировка n-наибольшая на основе одного столбца и условно на втором столбце

Question

У меня есть этот набор данных:

    individual  cluster  choice  benchmark_probabilities
0      9710535        0       0                 0.008647
1      9710535        2       0                 0.012558
2      9710535        2       0                 0.013894
3      9710535        1       0                 0.030648
4      9710535        1       0                 0.020298
5      9710535        1       0                 0.021444
6      9710535        1       0                 0.014804
7      9710535        5       0                 0.163837
8      9710535        5       0                 0.085191
9      9710535        2       0                 0.013272
10     9710535        2       0                 0.014684
11     9710535        2       0                 0.006987
12     9710535        2       0                 0.007387
13     9710535        2       0                 0.008940
14     9710535        3       0                 0.027746
15     9710535        3       0                 0.017345
16     9710535        3       0                 0.015545
17     9710535        4       0                 0.007449
18     9710535        3       0                 0.013382
19     9710535        4       0                 0.011559
20     9710535        3       0                 0.013091
21     9710535        4       0                 0.006438
22     9710535        4       0                 0.006089
23     9710535        4       0                 0.007768
24     9710535        4       0                 0.007348
25     9710535        2       0                 0.001479
26     9710535        5       0                 0.054764
27     9710535        5       0                 0.065420
28     9710535        5       0                 0.098600
29     9710535        5       0                 0.067577
30     9710535        6       0                 0.002158
31     9710535        6       0                 0.002041
32     9710535        6       0                 0.001694
33     9710535        6       0                 0.001602
34     9710535        7       0                 0.010075
35     9710535        7       0                 0.008076
36     9710535        7       0                 0.004485
37     9710535        7       0                 0.009090
38     9710535        7       0                 0.005834
39     9710535        5       0                 0.018973
40     9710535        7       0                 0.014945
41     9710535        7       0                 0.007159
42     9710535        6       0                 0.001624
43     9710535        6       0                 0.001535
44     9710535        5       0                 0.048068
45     9710535        7       0                 0.003548
46     9710540        0       1                 0.018614
47     9710540        0       0                 0.006515
48     9710540        0       0                 0.004040
49     9710540        1       0                 0.005489

, что я хочу сделать, это следующее:

1. Сгруппировать по individual, cluster и выбрать топ 1 каждой группы на основе benchmark_probabilities
2. Выберите 5 лучших результатов на основе individual
3. Если у одного individual меньше 5 уникальных cluster, заполните оставшееся пространство на основе на benchmark_probabilities без учета cluster.

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

    individual  cluster  choice  benchmark_probabilities
0     9710535        1       0                 0.030648
1     9710535        5       0                 0.163837
2     9710535        3       0                 0.027746
3     9710535        8       0                 0.015682
4     9710535       11       1                 0.050787
5     9710540        0       0                 0.004040
6     9710540        1       0                 0.005489
7     9710540        0       0                 0.006515
8     9710540        0       1                 0.018614

Я сделал следующее, что дает мне 1-й и 2-й этап, но не третий:

data.groupby(["individual", "cluster"])["benchmark_probabilities"].nlargest(1).groupby("individual").nlargest(5)

Но результат не тот, который я хочу, и он также выглядит некрасиво:

individual  individual  cluster     
9710535     9710535     5        7      0.163837
                        11       75     0.050787
                        1        3      0.030648
                        3        14     0.027746
                        8        49     0.015682
9710540     9710540     0        98     0.018614
                        1        101    0.005489

Любая помощь будет оценена

Answer 1

Я думаю, вам нужно DataFrame.sort_values с GroupBy.head вместо nlargest, потому что избегайте потерянного столбца choice, а также улучшите производительность:

df0 = (data.groupby(["individual", "cluster"])["benchmark_probabilities"].nlargest(1)
           .groupby("individual").nlargest(5))
print (df0)
individual  individual  cluster    
9710535     9710535     5        7     0.163837
                        1        3     0.030648
                        3        14    0.027746
                        7        40    0.014945
                        2        10    0.014684
9710540     9710540     0        46    0.018614
                        1        49    0.005489
Name: benchmark_probabilities, dtype: float64

df1 = (data.sort_values(['individual','cluster','benchmark_probabilities'],
                         ascending=[True, True, False])
           .groupby(["individual", "cluster"]).head(1)
           .sort_values(['individual','benchmark_probabilities'], 
                        ascending=[True, False])
           .groupby("individual").head(5))
print (df1)
    individual  cluster  choice  benchmark_probabilities
7      9710535        5       0                 0.163837
3      9710535        1       0                 0.030648
14     9710535        3       0                 0.027746
40     9710535        7       0                 0.014945
10     9710535        2       0                 0.014684
46     9710540        0       1                 0.018614
49     9710540        1       0                 0.005489

Затем отфильтруйте только строки из оригинала не в df1 и не выполните сортировку:

df2 = (data[~data.index.isin(df1.index)]
           .sort_values(['individual','benchmark_probabilities'], 
                        ascending=[True, False])
           )
#print (df2)

Добавил t df1 и получите значения top5 по head:

df = (pd.concat([df1, df2])
        .groupby('individual').head(5)
        .sort_values('individual'))
print (df)
    individual  cluster  choice  benchmark_probabilities
7      9710535        5       0                 0.163837
3      9710535        1       0                 0.030648
14     9710535        3       0                 0.027746
40     9710535        7       0                 0.014945
10     9710535        2       0                 0.014684
46     9710540        0       1                 0.018614
49     9710540        1       0                 0.005489
47     9710540        0       0                 0.006515
48     9710540        0       0                 0.004040