使加权随机状态函数可扩展

时间:2018-07-08 18:27:59

标签: python function random

我有一个加权随机状态函数,该函数接受有限数量的相对权重,并基于伪随机选择输出整数。但是,我想使该函数可扩展以处理n个状态/权重。用* args作为输入来重写此函数的一种优雅方法是什么?

编辑:由于这些权重是彼此相对的,所以对我来说棘手的部分是思考使elif逻辑可扩展的最佳方法。

def weighted_random(weight1, weight2, weight3, weight4, weight5, weight6, weight7, weight8):
    totalWeight = weight1 + weight2 + weight3 + weight4 + weight5 + weight6 + weight7 + weight8
    randomInt = random.randint(1, totalWeight)

    if randomInt <= weight1:
        return 0
    elif randomInt > weight1 and randomInt <= (weight1 + weight2):
        return 1
    elif randomInt > (weight1 + weight2) and randomInt <= (weight1 + weight2 + weight3):
        return 2
    elif randomInt > (weight1 + weight2 + weight3) and randomInt <= (weight1 + weight2 + weight3 + weight4):
        return 3
    elif randomInt > (weight1 + weight2 + weight3 + weight4) and randomInt <= (weight1 + weight2 + weight3 + weight4 + weight5):
        return 4
    elif randomInt > (weight1 + weight2 + weight3 + weight4 + weight5) and randomInt <= (weight1 + weight2 + weight3 + weight4 + weight5 + weight6):
        return 5
    elif randomInt > (weight1 + weight2 + weight3 + weight4 + weight5 + weight6) and randomInt <= (weight1 + weight2 + weight3 + weight4 + weight5 + weight6 + weight7):
        return 6
    elif randomInt > (weight1 + weight2 + weight3 + weight4 + weight5 + weight6 + weight7):
        return 7
    else:
        return("error")

2 个答案:

答案 0 :(得分:2)

您可以使用*args来获取所有参数作为列表。然后,您可以遍历它们以找出哪个randomInt被命中了。

def weighted_random(*weights):
    totalWeight = sum(weights)
    randomInt = random.randint(1, totalWeight)

    for i, weight in enumerate(weights):
        if randomInt <= weight:
            return i
        randomInt -= weight

    return "error"

如果这太易读或O(n)太快,请尝试以下操作:

def weighted_random(*weights):
    randomInt = random.randint(1, sum(weights))
    return next((
        i for i in range(len)
        if sum(weights[:i+1]) < randomInt <= sum(weights[:i+2])
    ), "error")

答案 1 :(得分:0)

如果您要创建一个使用累积权重的函数,则可以按照以下步骤做一些事情(假设Python 3.6 +):

import random 
import itertools as it

def wran(*weights):
    li=random.choices(list(range(len(weights))),cum_weights=list(it.accumulate(weights)))
    return li[0]

因此,假设您希望[0,1,2,3,4,5]的权重为[50,5,5,5,15,20](由于它们的总和为100,因此可以看作是百分比),您可以这样做:

from collections import Counter
>>> t=100000 
>>> [(k,'{:.6}%'.format(float(v)/t*100)) for k,v in sorted(Counter(wran(50,5,5,5,15,20) for _ in range(t)).items())]
[(0, '49.926%'), (1, '4.922%'), (2, '5.015%'), (3, '5.067%'), (4, '15.125%'), (5, '19.945%')]

接近假定的百分比。


如果要在Python 3.6+以外的其他Python上执行此操作,则可以生成列表的累加和的元组,然后使用bisect来获取该位置适合随机数的索引元组列表:

import bisect
def wran(*weights):
    totalWeight = sum(weights)
    randomInt = random.randint(1, totalWeight)
    li=[(sum(weights[0:i]),sum(weights[0:i+1])) for i in range(len(weights))]
    return bisect.bisect_left(li,(randomInt,))-1

对此进行测试:

>>> [(k,'{:.6}%'.format(float(v)/t*100)) for k,v in sorted(Counter(wran(50,5,5,5,20,5,5,5) for _ in range(t)).items())]
[(0, '49.943%'), (1, '4.979%'), (2, '5.048%'), (3, '4.968%'), (4, '20.06%'), (5, '5.059%'), (6, '4.954%'), (7, '4.989%')]

与最初发布的功能相比:

>>> [(k,'{:.6}%'.format(float(v)/t*100)) for k,v in sorted(Counter(weighted_random(50,5,5,5,20,5,5,5) for _ in range(t)).items())]
[(0, '49.986%'), (1, '5.203%'), (2, '4.962%'), (3, '4.998%'), (4, '19.977%'), (5, '5.004%'), (6, '4.986%'), (7, '4.884%')]