在Python中使用NumPy/Scipy高效均匀采样整数

Question

我遇到了一个问题，具体来说是根据一次随机的抛硬币结果，我需要从一个字符串中随机选一个起始位置。如果这个随机位置的选择是均匀的，我想到了两种方法来实现：一种是使用numpy.random里的多项式分布，另一种是用Python标准库里的简单randint函数。我是这样测试的：

from numpy import *
from numpy.random import multinomial
from random import randint
import time

def use_multinomial(length, num_points):
    probs = ones(length)/float(length)
    for n in range(num_points):
    result = multinomial(1, probs)

def use_rand(length, num_points):
    for n in range(num_points):
    rand(1, length)

def main():
    length = 1700
    num_points = 50000

    t1 = time.time()
    use_multinomial(length, num_points)
    t2 = time.time()
    print "Multinomial took: %s seconds" %(t2 - t1)

    t1 = time.time()
    use_rand(length, num_points)
    t2 = time.time()
    print "Rand took: %s seconds" %(t2 - t1)    

if __name__ == '__main__':
    main()

测试的结果是：

使用多项式分布耗时：6.58072400093秒
使用rand耗时：2.35189199448秒

看起来randint的速度更快，但我觉得还是挺慢的。有没有什么方法可以用numpy或scipy让这个过程快得多呢？

谢谢。

Answer 1

我把你的代码改了一下，让它能真正返回值（而且用了 randint 代替 rand，这不是你想要的吗？）像这样...

def use_multinomial(length, num_points):
    probs = ones(length)/float(length)
    return multinomial(1, probs, num_points)

def use_rand(length, num_points):
    return [randint(1,length) for _ in range(num_points)]

然后我试着写了自己的版本，使用 numpy.random.randint 来生成一个包含随机点的 numpy 数组：

def use_np_randint(length, num_point):
    return nprandint(1, length, num_points)

结果是：

Multinomial took: 13.6279997826 seconds
Rand took: 0.185000181198 seconds
NP randint took: 0.00100016593933 seconds

多项式的方法明显比较慢，但这真的是你想要的吗？我记得你说过你想要一个均匀分布？用 numpy 的 randint 显然是最快的选择。