按一列对每个组进行排序，并根据每个组获取顶部元素

Question

我有一个带有以下列的 pandas 数据框：

month	userID	probability	top_users_that_month
9	A	5.231	3
9	B	0.963	3
9	E	7.445	3
9	D	9.343	3
9	C	3.231	3
10	X	0.963	1
10	B	7.445	1
10	A	2.343	1
13	Y	1.963	2
13	B	9.445	2
13	D	0.343	2

我想创建一个标志列，其值为 1 或 0，具体取决于用户 ID 是否在每个月的 top_users 中（根据他们的probability 列）。请注意，每个月的热门用户数由top_users_that_month 列表示。

所需的输出如下：

month	userID	probability	top_users_that_month	flag_among_top_users
9	A	5.231	3	1
9	B	0.963	3	0
9	E	7.445	3	1
9	D	9.343	3	1
9	C	3.231	3	0
10	X	0.963	1	0
10	B	7.445	1	1
10	A	2.343	1	0
11	Y	1.963	2	1
11	B	9.445	2	1
11	D	0.343	2	0

我尝试使用groupby 和apply 创建列，但没有得到我想要的结果。

Answer 1

treshold = 5
df['flag_among_top_users'] = (df['probability'] > treshold).astype(int)

输出：

	month	userID	probability	top_users_that_month	flag_among_top_users
0	9	A	5.231	3	1
1	9	B	0.963	3	0
2	9	E	7.445	3	1
3	9	D	9.343	3	1
4	9	C	3.231	3	0
5	10	X	0.963	1	0
6	10	B	7.445	1	1
7	10	A	2.343	1	0
8	13	Y	1.963	2	0
9	13	B	9.445	2	1
10	13	D	0.343	2	0

Answer 2

由于每个组要标记的用户数量不同，最简单的方法是使用自定义函数，使用 groupby 和 apply。

例如：

def get_top_users(df):
    top_users = df['top_users_that_month'].iloc[0]
    idx = df.sort_values('probability', ascending=False).head(top_users).index
    df['flag_among_top_users'] = 0
    df.loc[idx, 'flag_among_top_users'] = 1
    return df
    
df.groupby('month').apply(get_top_users)

结果：

    month userID  probability  top_users_that_month  flag_among_top_users
0       9      A        5.231                     3                     1
1       9      B        0.963                     3                     0
2       9      E        7.445                     3                     1
3       9      D        9.343                     3                     1
4       9      C        3.231                     3                     0
5      10      X        0.963                     1                     0
6      10      B        7.445                     1                     1
7      10      A        2.343                     1                     0
8      13      Y        1.963                     2                     1
9      13      B        9.445                     2                     1
10     13      D        0.343                     2                     0

也可以使用nlargest 而不是在get_top_users 内部进行排序，但如果有两个相同的probability 值可能会出现问题。