Simulation for conditional probabilty problem in python

2018-11-01

An answer to this question on Stack Overflow.

Question

I am trying to simulate a simple conditional probability problem. You hae two boxes. If you open A you have a 50% change of wining the prize, If you open B you have a 75% chance of winning. With some simple (bad) python I have tired But the appending doesn't work. Any thoughts on a neater way of doing this?

import random
import numpy as np
def liveORdie(prob):
    #Takes an argument of the probability of survival
    live = 0
    for i in range(100):
        if random.random() <= prob*1.0:
            live =1
    return live
def simulate(n):
    trials = np.array([0])
    for i in range(n):
        if random.random() <= 0.5:
            np.append(trials,liveORdie(0.5))
            print(trials)
        else:
            np.append(trials,liveORdie(0.75))
    return(sum(trials)/n)
   
simulate(10)

Answer

You could make the code tighter by using list comprehensions and numpy's array operations, like so:

import random
import numpy as np
def LiveOrDie():
    prob = 0.5 if random.random()<=0.5 else 0.75
    return np.sum(np.random.random(100)<=prob)
def simulate(n):
    trials = [LiveOrDie() for x in range(n)]
    return(sum(trials)/n)
Simulate(10)