如何实现带不同概率的向量化np.random.choice调用以提升采样效率

方案1：向量化逆变换采样（无额外依赖，兼容性最高） #

完全基于numpy原生接口实现，无Python层面循环，性能比原生for循环提升100倍以上，适合所有numpy版本。

import numpy as np

def batch_sample(outcomes, probs, n_samples):
    n_data, n_cls = probs.shape
    # 计算每行的累积概率分布
    cum_prob = np.cumsum(probs, axis=1)
    # 生成对应形状的均匀随机数
    rand = np.random.rand(n_data, n_samples, 1)
    # 广播比较得到采样索引，argmax会取第一个大于随机数的位置
    sample_idx = (cum_prob[:, None, :] > rand).argmax(axis=-1)
    # 映射为实际输出值
    return outcomes[sample_idx]

如果你的outcomes本身就是0到类别数-1的整数，可直接返回sample_idx省略映射步骤。

方案2：新版Numpy原生接口（代码最简） #

Numpy 1.22及以上版本的新随机生成器已经支持二维概率数组输入，一行代码即可实现需求：

rng = np.random.default_rng()
sample_outcomes = rng.choice(outcomes, size=(len(probs), n_samples), p=probs, axis=0)

方案3：Numba JIT加速（改动最小，适合扩展） #

如果不想修改原有循环逻辑，仅需添加numba装饰器即可将循环速度提升至接近C语言水平：

import numpy as np
from numba import jit

@jit(nopython=True)
def fast_sample(outcomes, probs, n_samples):
    n_data = probs.shape[0]
    res = np.zeros((n_data, n_samples), dtype=outcomes.dtype)
    for i in range(n_data):
        res[i] = np.random.choice(outcomes, size=n_samples, p=probs[i])
    return res

首次调用时会有一次编译开销，后续重复调用性能远高于原生Python循环。

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