python - 基于一列的 MAX 值和有条件地另一列的 topN 值进行分组和过滤

Question

我有以下数据:

    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

我想根据每个 groupID 的 probA 列的最大值(即 groupID = 5)过滤每个individual 的数据 for 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

如果有任何帮助，我将不胜感激!

Answer 1

我认为您需要首先通过 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 ，并通过 GroupBy.head 获得前 5 名最后按另外两列排序:

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