Pandas按月份分组取指定数量概率头部用户并生成0/1标记列

方法1：groupby + rank 实现（性能最优，无遍历开销） #

该方法无需自定义函数遍历分组，适合数据量较大的场景：

import pandas as pd

# 打标逻辑实现
df['flag_among_top_users'] = (
    # 按month分组，对probability做降序排名
    df.groupby('month')['probability']
    .rank(ascending=False, method='first')
    # 判断排名是否小于等于当月指定的top用户数量
    .le(df['top_users_that_month'])
    # 布尔值转0/1
    .astype(int)
)

参数说明：method='first'表示遇到相同probability时，按数据出现的先后顺序排名，避免排名重复导致打标数量超出阈值。如果业务要求相同概率的用户全部标记为1，可将参数替换为method='min'。

方法2：groupby.apply 正确实现（适配自定义逻辑扩展场景） #

如果之前使用apply未得到预期结果，大多是分组处理后索引未对齐、排序逻辑有误导致，正确写法如下：

def mark_top_user_per_month(group):
    # 取当前月指定的top用户数量（同month下该字段值统一）
    top_n = group['top_users_that_month'].iloc[0]
    # 按probability降序排序，取前top_n个用户的索引
    top_user_index = group.sort_values('probability', ascending=False).head(top_n).index
    # 初始化标记为0，给头部用户赋值为1
    group['flag_among_top_users'] = 0
    group.loc[top_user_index, 'flag_among_top_users'] = 1
    return group

# group_keys=False避免生成多级索引，保证返回结果和原表索引对齐
df = df.groupby('month', group_keys=False).apply(mark_top_user_per_month)

提示：如果业务中同一month下top_users_that_month存在不同取值，需要先统一该字段的月度取值后再执行上述逻辑

内容的提问来源于stack exchange，提问作者Javier Monsalve