如何对Pandas DataFrame中的循环计算进行向量化？

我需要对包含prob、id、time字段的Pandas DataFrame执行以下for循环计算，想改成向量化实现以提升效率：

确定性示例输入生成代码 #

import pandas as pd
import random

random.seed(0)
probs = [i / 10 for i in range(11)]
ids = [i for i in range(5)]
rowcounts = [i for i in range(2, 6)]
times = [i for i in range(14)]

example = []
for idn in ids:
    rowcount = random.choice(rowcounts)
    pvals = sorted(random.sample(probs, rowcount))
    tvals = sorted(random.sample(times, rowcount))
    for i in range(rowcount):
        example.append([idn, pvals[i], tvals[i]])
df = pd.DataFrame(example, columns=['id', 'prob', 'time'])

原循环计算逻辑 #

import numpy as np

data = []
for cutoff in np.sort(df['prob'].unique()):
    # 筛选prob >= cutoff的记录
    sub = df[df['prob'] >= cutoff]
    # 每个id取prob最小的记录
    subsub = sub.loc[sub.groupby('id')['prob'].idxmin()]
    # 计算time字段的统计量
    times = subsub['time'].values
    data.append([cutoff, np.quantile(times, 0.5), np.quantile(times, 0.9), (times <= 7).sum()/len(times)])
metrics = pd.DataFrame(data)

向量化解决方案 #

核心思路是利用pd.merge_asof实现批量匹配每个cutoff下各id的有效time，再通过分组统计直接计算所需指标，完全避免循环：

import pandas as pd
import numpy as np

# 1. 准备排序后的cutoff列表
cutoff_df = pd.DataFrame({'cutoff': np.sort(df['prob'].unique())})

# 2. 生成所有id与cutoff的组合
id_cutoff = cutoff_df.assign(key=1).merge(
    pd.DataFrame({'id': df['id'].unique(), 'key': 1}),
    on='key'
).drop('key', axis=1)

# 3. 对原数据按id和prob升序排序，确保匹配逻辑正确
sorted_df = df.sort_values(['id', 'prob'])
id_cutoff = id_cutoff.sort_values(['id', 'cutoff'])

# 4. 使用merge_asof批量匹配每个(id, cutoff)对应的最小prob>=cutoff的time
merged = pd.merge_asof(
    id_cutoff,
    sorted_df,
    left_on='cutoff',
    right_on='prob',
    by='id',
    direction='forward'  # 找第一个>=cutoff的prob对应的行
)

# 5. 按cutoff分组计算统计量，与原循环逻辑完全一致
metrics_vectorized = merged.groupby('cutoff')['time'].agg(
    median=lambda x: np.quantile(x.dropna(), 0.5),
    q90=lambda x: np.quantile(x.dropna(), 0.9),
    ratio=lambda x: (x.dropna() <= 7).sum() / len(x.dropna())
).reset_index()

# 验证结果一致性（可选）
print("原循环结果与向量化结果是否一致：", metrics.round(6).equals(metrics_vectorized.round(6)))

逻辑说明 #

1. 生成id-cutoff组合：覆盖所有需要计算的cutoff和id的配对，确保每个cutoff下的所有id都被考虑。
2. merge_asof匹配：direction='forward'会为每个(id, cutoff)找到该id中第一个prob>=cutoff的行，正好对应原循环中"每个id取prob最小的符合条件的记录"的逻辑。
3. 分组统计：直接对每个cutoff分组，计算time的中位数、90分位数和占比，无需重复筛选和分组操作。

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