如何高效筛选Pandas DataFrame每行各分类Top N列并置零落选值？

数据集背景 #

假设我们有一个记录选手每日酷感得分的Pandas DataFrame，生成代码如下：

import numpy as np
import pandas as pd

rng = np.random.default_rng()
dates = pd.date_range('2024-08-01', '2024-08-07')
contestants = ['Alligator', 'Beryl', 'Chupacabra', 'Dandelion', 'Eggplant', 'Feldspar']
coolness_score = pd.DataFrame(rng.random((len(dates), len(contestants))), index=dates, columns=contestants)

样例数据输出：

Alligator     Beryl  Chupacabra  Dandelion  Eggplant  Feldspar
2024-08-01   0.213901  0.952705    0.801651   0.511080  0.662109  0.486296
2024-08-02   0.495700  0.660502    0.379900   0.778438  0.038616  0.214174
2024-08-03   0.639337  0.036226    0.811501   0.281915  0.101850  0.437146
2024-08-04   0.238590  0.686965    0.357087   0.810922  0.907803  0.370247
2024-08-05   0.712564  0.800191    0.040616   0.503644  0.354333  0.742269
2024-08-06   0.916343  0.299557    0.405399   0.851161  0.336570  0.246618
2024-08-07   0.047052  0.645420    0.823397   0.198483  0.368888  0.168188

每位选手对应固定类别，且各类别有入选数量限制：

category_mapping = {
    'Alligator': 'Animal',
    'Beryl': 'Mineral',
    'Chupacabra': 'Animal',
    'Dandelion': 'Vegetable',
    'Eggplant': 'Vegetable',
    'Feldspar': 'Mineral'
}

category_limits = {
    'Animal': 1,
    'Vegetable': 2,
    'Mineral': 1
}

需求说明 #

需要实现三类筛选逻辑，落选选手的得分需置零：

• 场景1：每个类别选取单条最高分
• 场景2：每个类别选取固定数量N的最高分
• 场景3：按category_limits定义的数量选取每个类别的最高分

实际数据集为250行×15000列，包含175个类别，需高频用于蒙特卡洛模拟，必须避免逐行迭代，采用高效矢量化操作。

高效实现方案 #

通用预处理：构建类别分组映射 #

先将选手按类别分组，生成类别-选手列表的映射：

from collections import defaultdict

category_groups = defaultdict(list)
for contestant, cat in category_mapping.items():
    category_groups[cat].append(contestant)

场景1：每个类别取单条最高分 #

利用groupby+transform计算类别内每日最高分，再通过布尔索引保留最高分，其余置零：

# 将类别映射转为Series，方便按列分组
cat_series = pd.Series(category_mapping)

# 计算每个类别内的每日最高分
max_per_cat = coolness_score.groupby(cat_series, axis=1).transform('max')

# 保留最高分，其余置零
result_scenario1 = coolness_score.where(coolness_score == max_per_cat, 0)

场景2：每个类别取固定数量N的最高分 #

假设N=2，通过rank函数对类别内得分降序排名，保留前N名得分：

N = 2

# 按类别分组，对每行得分降序排名（相同得分取相同排名）
rank_per_cat = coolness_score.groupby(cat_series, axis=1).rank(method='min', ascending=False)

# 保留排名<=N的得分，其余置零
result_scenario2 = coolness_score.where(rank_per_cat <= N, 0)

场景3：按自定义类别限制取最高分 #

结合category_limits，对每个类别单独处理后合并结果，全程采用矢量化操作：

# 初始化全零结果DataFrame
result_scenario3 = pd.DataFrame(0, index=coolness_score.index, columns=coolness_score.columns)

for cat, limit in category_limits.items():
    # 获取当前类别的选手列
    cols = category_groups[cat]
    # 对每行得分降序排名
    ranks = coolness_score[cols].rank(method='min', ascending=False, axis=1)
    # 保留排名<=limit的得分，其余置零
    result_scenario3[cols] = coolness_score[cols].where(ranks <= limit, 0)

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