求助：基于NumPy向量化优化置信集与类别覆盖率计算代码

优化后的完整代码 #

import numpy as np

def get_confidence_sets(y_pred, q_yhat):
    # 对每行降序排序，同时获取原始类索引
    sorted_vals, sorted_indices = np.sort(y_pred, axis=1)[:, ::-1], np.argsort(y_pred, axis=1)[:, ::-1]
    # 计算每行的累积和
    cumsum_vals = np.cumsum(sorted_vals, axis=1)
    # 找到每行第一个累积和 >= q_yhat的位置（argmax返回第一个True的索引）
    a_indices = np.argmax(cumsum_vals >= q_yhat, axis=1)
    # 生成掩码：每个样本保留前a_indices[i]+1个类
    mask = np.arange(y_pred.shape[1]) <= a_indices[:, np.newaxis]
    # 提取每个样本的置信集合
    conf_sets = [sorted_indices[i][mask[i]].tolist() for i in range(y_pred.shape[0])]
    return conf_sets

def get_class_coverage(true_label, conf_sets, nb_class):
    # 将conf_sets转换为nb_samples x nb_classes的布尔掩码
    nb_samples = len(conf_sets)
    coverage_mask = np.zeros((nb_samples, nb_class), dtype=bool)
    for i, classes in enumerate(conf_sets):
        coverage_mask[i, classes] = True
    # 计算d2：每个类的真实样本数
    d2 = np.bincount(true_label, minlength=nb_class)
    # 计算d1：每个类中真实标签在置信集合中的样本数
    is_covered = coverage_mask[np.arange(nb_samples), true_label]
    d1 = np.bincount(true_label[is_covered], minlength=nb_class)
    return d1, d2

关键优化细节 #

1. get_confidence_sets的优化

• 移除嵌套循环：原代码内层循环通过np.where(y_pred==...)反复查找类索引，会重复遍历整个矩阵，效率极低。改用np.argsort直接获取每行降序排序后的原始类索引，一步到位拿到所有排序后的类位置。
• 向量化计算累积和与截断位置：用np.cumsum(axis=1)对所有行同时计算累积和，再用np.argmax定位每行第一个满足条件的索引（因累积和递增，第一个True即为目标位置），比np.where(...)[0][0]更高效。
• 掩码批量筛选：生成布尔掩码一次性标记所有样本需要保留的类，避免逐样本循环筛选。

2. get_class_coverage的优化

• 转换为布尔掩码：将conf_sets的列表结构转为二维布尔数组，利用NumPy向量化操作快速判断每个样本的真实标签是否在置信集合中。
• 用np.bincount替代循环计数：原代码循环统计类样本数和覆盖数，改用np.bincount直接对真实标签数组做频次统计，速度提升数个数量级。
• 消除线性查找：原代码中y_val[i] in conf_sets[i]是O(k)线性查找（k为置信集合大小），转为掩码后直接通过索引访问，是O(1)操作。

额外性能建议 #

• 如果业务允许，conf_sets可以直接保留布尔掩码或二维索引数组，完全避免列表遍历，进一步提速。
• 若数据精度允许，将y_pred转为float32类型，GPU运算时会比float64更快。

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