如何用Pandas向量式操作按组生成对应累计最大prob的accuracy列？

原始数据与需求 #

原始DataFrame代码 #

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "realization_id": np.repeat([0, 1], 6),
    "sample_size": np.tile([0, 1, 2], 4),
    "num_obs": np.tile(np.repeat([25, 100], 3), 2),
    "accuracy": [0.8, 0.7, 0.8, 0.6, 0.7, 0.5, 0.6, 0.7, 0.8, 0.7, 0.9, 0.7],
    "prob": [0.94, 0.96, 0.95, 0.98, 0.93, 0.92, 0.90, 0.92, 0.95, 0.9, 0.91, 0.92]
})

需求说明 #

需要生成desired_accuracy列，规则如下：

• 以(realization_id, num_obs)为分组依据，每组内规则独立重置
• 每行的desired_accuracy值，对应到当前行为止，组内**首次达到最大prob**的行的accuracy值
• 后续行如果prob未超过当前累积最大值，则保持之前的desired_accuracy不变

目标输出列示例：

df["desired_accuracy"] = [0.8, 0.7, 0.7, 0.6, 0.6, 0.6, 0.6, 0.7, 0.8, 0.7, 0.9, 0.7]

向量式实现方案 #

通过groupby结合transform实现全程向量式操作，避免循环：

# 生成desired_accuracy列
df["desired_accuracy"] = df.groupby(["realization_id", "num_obs"]).transform(
    lambda g: g["accuracy"].where(g["prob"] == g["prob"].cummax()).ffill()
)

代码解释 #

1. 分组隔离：以(realization_id, num_obs)分组，确保每组规则独立生效
2. 累积最大值标记：g["prob"].cummax()计算组内到当前行的累积最大prob，通过相等判断生成掩码，标记出首次刷新最大值的行
3. 保留对应accuracy：用where仅保留掩码为True的行的accuracy，其余行设为NaN
4. 向前填充：ffill()填充NaN，让未刷新最大值的行继承之前的accuracy值

执行后，df["desired_accuracy"]将完全匹配目标输出。

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