Search code examples
pythonrandomprobability

How to choose keys from a python dictionary based on weighted probability?


I have a Python dictionary where keys represent some item and values represent some (normalized) weighting for said item. For example:

d = {'a': 0.0625, 'c': 0.625, 'b': 0.3125}
# Note that sum([v for k,v in d.iteritems()]) == 1 for all `d`

Given this correlation of items to weights, how can I choose a key from d such that 6.25% of the time the result is 'a', 32.25% of the time the result is 'b', and 62.5% of the result is 'c'?


Solution

  • def weighted_random_by_dct(dct):
        rand_val = random.random()
        total = 0
        for k, v in dct.items():
            total += v
            if rand_val <= total:
                return k
        assert False, 'unreachable'
    

    Should do the trick. Goes through each key and keeps a running sum and if the random value (between 0 and 1) falls in the slot it returns that key