Search code examples
pythonlistprobability

Is there a neater way to implement: do this with probability p and that with probability (1-p)? (python)


I am trying to implement the following in Python: 'with probability p randomly choose a variable from a list l and with probability (1-p) choose a fixed variable w'. I have an implementation based on this question which seems to work, but I was just wondering if there is a neater way of doing this? The following is for probability 50%, but I want this to work for any probability p (to 2 decimals):

l = [1,2,3,4,5,6]

p = 0.5 # this is my probability p
r = 100*p # this transforms my probability p to an integer
q = 100*(1-p) # this transforms my probability (1-p) to an integer

v = random.choice(l) # a randomly chosen variable from a list l
w = 2 # a fixed variable

listt = []
listt.extend([v]*int(r)) # this adds a randomly chosen v to the list r times
listt.extend([w]*int(q)) # this adds w to the list q times

# I now have a list which has r elements v and q elements w

random.choice(listt)

Thanks!


Solution

  • random.random() makes for a simpliest generator against float-value probability.

    import random
    
    p = 0.2
    
    for _ in range(20):
        if random.random() < p:
            print('0') # probability p
        else:
            print('1') # probability 1-p
    

    And simpliest way to emulate specifically what you were doing in your question:

    import random
    p = 0.2
    
    v = 0
    w = 1
    
    
    for _ in range(20):
        print(random.choices([v, w], [p, 1-p])[0]) # [0] as choices returns a list and we only care about one value.
    

    random.choices accepts a weights parameter. if you pass p and 1-p for each element respectively, you will get elements chosen with correct probabilities.