使加权随机状态函数可拓

Question

我有一个加权随机状态函数，它接受有限个相对权重，并根据伪随机选择输出一个整数。但是，我想使函数可扩展以处理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")

Answer 1

可以使用*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")

Answer 2

如果您试图创建一个使用累积权重的函数，那么可以按照以下几行进行操作（假设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%')]