# Selecting random class from weighted class probability distribution

I have a solution that I've scratched out, but it's a tad messy:

First, let's say we have a random number from a uniform distribution we'll be using to select a random class from a class probability distribution:

r = 0.525325235325


And a list of classes, along with a corresponding list of values representing a class probability distribution for the four classes:

classes = ["w", "x", "y", "z"]

distribution = [0.1, 0.2, 0.3, 0.4]


Now, to get the cumulative probabilities:

b = [sum(distribution[0:x + 1]) for x in range(len(distribution))]
>>> b
[0.1, 0.30000000000000004, 0.6000000000000001, 1.0]


So far so good. Now, select a class by the weighted probability:

c = [y for y in b if y > r][0]
>>> c
0.6000000000000001


And to get the selected class:

selected = classes[b.index(c)]
>>> selected
'y'


Putting it all together:

r = 0.525325235325
classes = ["w", "x", "y", "z"]
distribution = [0.1, 0.2, 0.3, 0.4]

b = [sum(distribution[0:x + 1]) for x in range(len(distribution))]
selected = classes[b.index([y for y in b if y > r][0])]


This works well enough. The thing is, it will be called each time the class probability distribution is updated along with a complex code base – so efficiency matters.

Are there any ways I can make this code more efficient in terms of speed and/or memory usage (eliminating b without the expense of additional processing would be nice) in Python? Alternatively, is there an established algorithm for this problem that I've overlooked?

Using dictionaries is faster than lists. Your distributions can serve as keys to find the corresponding values instead of using indexed lists.

• That's a fair point, I'll look into it. The code base may make that tricky off the cuff, but nothing a little refactoring can't fix. Thanks : ) May 7, 2017 at 22:50