python之在 Python/NumPy 中加速多项式随机样本

我正在根据一组概率 probs 的多项式分布生成抽奖向量，其中每次抽奖都是 probs 中所选条目的索引:

import numpy as np 
def sample_mult(K, probs): 
    result = np.zeros(num_draws, dtype=np.int32) 
    for n in xrange(K): 
        draws = np.random.multinomial(1, probs) 
        result[n] = np.where(draws == 1)[0][0] 
    return result 

这可以加速吗？一遍又一遍地调用 np.random.multinomial 似乎效率很低(而且 np.where 也可能很慢。)

timeit 表示 最慢的运行时间比最快的运行时间长 6.72 倍。这可能意味着正在缓存中间结果 100000 次循环，3 次循环中的最佳循环:每次循环 18.9 µs

请您参考如下方法:

您可以将 size 选项与 np.random.multinomial 一起使用使用默认 size=1 输出多行随机样本，而不是仅输出一行，然后使用 .argmax(1)模拟np.where()[0][0]行为。

因此，我们会有一个矢量化的解决方案，就像这样 -

result = (np.random.multinomial(1,probs,size=K)==1).argmax(1)