python按概率生成随机数

Talk is cheap, show me the code.

random_demo.py代码如下:

# coding: utf-8
import random

def random_index(rate):
    """随机变量的概率函数"""
    #
    # 参数rate为list
    # 返回概率事件的下标索引
    start = 0
    index = 0
    randnum = random.randint(1, sum(rate))

    for index, scope in enumerate(rate):
        start += scope
        if randnum <= start:
            break
    return index

def main():
    arr = ['red', 'green', 'blue']
    rate = [45, 30, 25]

    red_times = 0
    green_times = 0
    blue_times = 0
    for i in xrange(10000):
        if arr[random_index(rate)] == 'red':
            red_times += 1
        if arr[random_index(rate)] == 'green':
            green_times += 1
        if arr[random_index(rate)] == 'blue':
            blue_times += 1

    print red_times, green_times, blue_times

if __name__ == '__main__':
    main()

打印结果:

$ python random_demo.py 
4507 3009 2537
$ python random_demo.py 
4586 3055 2485
$ python random_demo.py 
4483 2996 2506
$ python random_demo.py 
4586 3015 2480
$ python random_demo.py 
4513 3057 2415
$ python random_demo.py 
4494 3014 2456
$ python random_demo.py 
4502 2978 2505
$ python random_demo.py 
4475 3041 2528
$ python random_demo.py 
4561 3043 2520
$ python random_demo.py 
4558 2923 2474
$ python random_demo.py 
4536 2957 2535
$ python random_demo.py 
4526 2985 2576
$ python random_demo.py 
4532 2936 2547
$ python random_demo.py 
4420 2908 2447

可以看到结果, 抽中红绿蓝的次数非常接近它们的概率比例45:30:25, 其中random_index(rate)函数为核心代码.

OK, Enjoy it~