如何用Numpy实现逐行加权的random.choice批量采样？

你可以用以下几种高效的批量实现方式替代手动for循环：

1. NumPy 基于累积分布（CDF）的向量化实现 #

这是效率最高的方案，完全摆脱Python循环开销：

import numpy as np

def batch_weighted_sample(a, W):
    N, _ = W.shape
    # 计算每行权重的累积分布函数（CDF）
    cdf = W.cumsum(axis=1)
    # 生成N个[0,1)区间的均匀随机数
    u = np.random.rand(N, 1)
    # 找到每个随机数对应的抽样索引
    samples = (cdf < u).sum(axis=1)
    # 映射到目标数组a的元素
    return a[samples]

# 使用示例
N = 1000
T = 5
W = np.random.dirichlet(np.ones(T+1), size=N)  # 生成归一化权重矩阵
a = np.arange(T+1)
samples = batch_weighted_sample(a, W)

2. SciPy 离散分布批量抽样 #

利用scipy.stats.rv_discrete封装好的离散分布实现批量抽样：

from scipy.stats import rv_discrete
import numpy as np

def batch_scipy_sample(a, W):
    # 创建支持批量权重的离散分布对象
    dist = rv_discrete(values=(a, W.T))
    # 抽取对应N个样本
    return dist.rvs(size=W.shape[0])

# 使用示例
samples = batch_scipy_sample(a, W)

3. NumPy apply_along_axis 封装（不推荐大N场景） #

本质是隐式循环，效率接近手动for循环，仅适用于小数据量：

samples = np.apply_along_axis(lambda p: np.random.choice(a, p=p), axis=1, arr=W)

关键说明 #

• 基于CDF的NumPy方案：是性能最优的选择，所有操作都是向量化的numpy内部运算，适合大规模N的场景
• SciPy方案：代码更简洁，封装性强，底层实现同样高效
• apply_along_axis方案：代码最简洁但效率最低，仅作为备选

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