How to generate random int with given probability in python?

Question

I want to output a random flag (either 0,1 or 2) based on given probabilities.

eg

flag = gen_flag(p0,p1,p2)

where p0+p1+p2=1 and p0 , p1 , and p2 indicate the probabilities generating flag 0,1, and 2 respectively. So for example, gen_flag(0.8,0.1,0.1) would very likely output a zero.

How do I do that?

Answer 1

Use random.choices

from random import choices

flag = choices([0,1,2], [p0, p1, p2])[0]

p0 et al. don't have to sum to 1; they are normalized if they don't already.

choices always returns a list, even in the default case of choosing only 1 element from the set of possible flags.

Answer 2

@Chepner has the best answer, but choices wasn't added to random until python 3.6. For earlier versions, I would take the inputs and make a 100 item list with the potential number of 0s, 1s, and 2s respectively. Then use the random module to pick an item.

import random

gen_flag(p0, p1, p2):
    myList = [0]*int(p0*100) + [1]*int(p1*100) + [2]*int(p2*100)
    return myList[random.randint(0,99)]