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

Question

我正在根据一组概率 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

Answer 1

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

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

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