时间戳

首页 > 解决方案 > Python：numpy 的 random.choice() 的更快替代方案？

python - Python：numpy 的 random.choice() 的更快替代方案？

问题描述

我正在尝试对 0 到 999 之间的 1000 个数字进行采样，权重向量指示选择特定数字的概率：

import numpy as np
resampled_indices = np.random.choice(a = 1000, size = 1000, replace = True, p = weights)

不幸的是，这个过程必须在更大的 for 循环中运行数千次，这似乎np.random.choice是该过程的主要速度瓶颈。因此，我想知道是否有任何方法可以加快速度np.random.choice或使用产生相同结果的替代方法。

标签： pythonnumpyrandom

解决方案

您似乎可以通过使用统一采样然后使用“反转”累积分布来稍微快一点np.searchsorted：

# assume arbitrary probabilities
weights = np.random.randn(1000)**2
weights /= weights.sum()

def weighted_random(w, n):
    cumsum = np.cumsum(w)
    rdm_unif = np.random.rand(n)
    return np.searchsorted(cumsum, rdm_unif)

# first method
%timeit np.random.choice(a = 1000, size = 1000, replace = True, p = weights)
# 10000 loops, best of 3: 220 µs per loop

# proposed method
%timeit weighted_random(weights, n)
# 10000 loops, best of 3: 158 µs per loop

现在我们可以凭经验检查概率是否正确：

samples =np.empty((10000,1000),dtype=int)
for i in xrange(10000):
    samples[i,:] = weighted_random(weights)

empirical = 1. * np.bincount(samples.flatten()) / samples.size
((empirical - weights)**2).max()
# 3.5e-09

推荐阅读