如何在Pandas DataFrame中高效按组标记列最大值？（2000列场景）

针对你处理2000个draw列的场景，循环遍历每个列并重复执行groupby确实会带来巨大的性能开销——这是因为每次循环都要重新计算分组，2000次重复操作的累积成本很高。我们可以利用Pandas的向量化分组操作，一次性完成所有draw列的标记，效率会提升几个数量级。

核心思路 #

1. 分离数据中的非draw列和draw列，避免不必要的计算；
2. 仅对draw列执行一次分组操作，用transform生成与原DataFrame同形状的组内最大值矩阵；
3. 通过元素级比较，直接将等于组内最大值的位置标记为1，其余为0；
4. 合并处理后的draw列与原非draw列，得到最终结果。

完整实现代码 #

import pandas as pd
import numpy as np

# 生成模拟数据（含2000个draw列）
def simulateDraw(df, n=2000):
    return pd.concat([
        df, 
        df.apply(lambda row: pd.Series(np.random.beta(row.C, row.C, size=n)), axis=1).add_prefix('draw')
    ], axis=1)

# 初始化示例数据
data = {'A':[11111, 11111, 33333,11111], 'B':[101, 101, 102, 101],'C':[1,2,3,4]}
df = pd.DataFrame(data)
df = simulateDraw(df, n=2000)

# 分离非draw列和draw列
non_draw_cols = [col for col in df if not col.startswith('draw')]
draw_cols = [col for col in df if col.startswith('draw')]

# 高效标记最大值行
# 1. 计算每个组内draw列的最大值（transform保持原形状）
group_max = df.groupby(['A', 'B'])[draw_cols].transform('max')
# 2. 比较原数据与组内最大值，相等则为1，否则为0
df[draw_cols] = (df[draw_cols] == group_max).astype(int)

# 查看结果示例（取前3个draw列）
print(df[non_draw_cols + draw_cols[:3]])

为什么这个方法更快？ #

• 避免重复分组：原循环中每个draw列都要执行一次groupby，而新方法只需要分组一次，大幅减少了分组计算的开销；
• 向量化操作：所有列的比较和赋值都是一次性完成的，利用了Pandas底层的C级优化，比Python循环快得多；
• 内存效率：transform生成的最大值矩阵与原数据同形状，不需要额外的索引查找（比如原方法中的idxmax和isin操作），减少了内存占用和索引匹配的成本。

示例输出（针对初始小数据） #

假设我们用初始的3个draw列测试，处理后结果如下：

A    B  C  draw0  draw1  draw2
0  11111  101  1      0      1      0
1  11111  101  2      1      0      0
2  33333  102  3      1      1      1
3  11111  101  4      0      0      1

可以看到，每个[A,B]组内的draw列最大值行都被正确标记为1，其余为0。

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