algorithm - Python Weighted Random

26
votes

I need to return different values based on a weighted round-robin such that 1 in 20 gets A, 1 in 20 gets B, and the rest go to C.

So:

A => 5%
B => 5%
C => 90%

Here's a basic version that appears to work:

import random

x = random.randint(1, 100)

if x <= 5:
    return 'A'
elif x > 5 and x <= 10:
    return 'B'
else:
    return 'C'

Is this algorithm correct? If so, can it be improved?

4
pythonalgorithmround-robin
You could use random.randint(1,20) for your case. - Akavall
@Akavall - How so? (1,20) would only return only let me eval if only A or B fell into a 5% range, but not both - right? - doremi
Your random integer can take value 1 to 20, if random integer is 1, you return A (5% chance), if random integer is 2, you return B (5% chance), if random integer is anything else you return C (90% likelihood). Am I missing something? - Akavall
6 in 1, half a dozen in the other. Your logic works now that I think about it. - doremi
In case you wanted to research this in the more general case, the concept you are referring to is "generating random variates using the inverse cdf method." - Joel Cornett

4 Answers

56
votes

Your algorithm is correct, how about something more elegant:

import random
my_list = ['A'] * 5 + ['B'] * 5 + ['C'] * 90
random.choice(my_list)
37
votes

that's fine. more generally, you can define something like:

from collections import Counter
from random import randint

def weighted_random(pairs):
    total = sum(pair[0] for pair in pairs)
    r = randint(1, total)
    for (weight, value) in pairs:
        r -= weight
        if r <= 0: return value

results = Counter(weighted_random([(1,'a'),(1,'b'),(18,'c')])
                  for _ in range(20000))
print(results)

which gives

Counter({'c': 17954, 'b': 1039, 'a': 1007})

which is as close to 18:1:1 as you can expect.

9
votes

If you want to use weighted random and not percentile random, you can make your own Randomizer class:

import random

class WeightedRandomizer:
    def __init__ (self, weights):
        self.__max = .0
        self.__weights = []
        for value, weight in weights.items ():
            self.__max += weight
            self.__weights.append ( (self.__max, value) )

    def random (self):
        r = random.random () * self.__max
        for ceil, value in self.__weights:
            if ceil > r: return value

w = {'A': 1.0, 'B': 1.0, 'C': 18.0}
#or w = {'A': 5, 'B': 5, 'C': 90}
#or w = {'A': 1.0/18, 'B': 1.0/18, 'C': 1.0}
#or or or

wr = WeightedRandomizer (w)

results = {'A': 0, 'B': 0, 'C': 0}
for i in range (10000):
    results [wr.random () ] += 1

print ('After 10000 rounds the distribution is:')
print (results)
0
votes

It seems correct since you are using a uniform random variable with independent draws the probability for each number will be 1/n (n=100).

You can easily verify your algorithm by running it say 1000 time and see the frequency for each letter.

Another algorithm you might consider is to generate an array with your letters given the frequency you want for each letter and only generate a single random number which is the index in the array

It will be less efficient in memory but should perform better

Edit:

To respond to @Joel Cornett comment, an example will be very similar to @jurgenreza but more memory efficient 

import random
data_list = ['A'] + ['B'] + ['C'] * 18
random.choice(data_list )