Probability puzzle and simulation

ball_quiz_simulation.py

#!/usr/bin/env python
"""
This script is for a Monte Carlo simulation of the following puzzle: you
first put a white ball in a bag containing another "original" ball with
unknown color, which can be black or white with probability of 1/2. The
question is what is the probability that the original ball was white if
you pick one of two balls in the bag randomly and the chosen ball is
white. As a matter of fact, this problem is equivalent to finding the
probability of choosing a white ball out of one black and two white
balls, which should be 2/3.

If you have any comment, please leave it at
http://serizawa.egloos.com/1848317
"""

import sys
from random import random as rand

#color = {0: 'white', 1: 'black'}
#choice = {0: 'original ball', 1: 'new ball'}

def main(Ntrials):
    original_ball = [int(rand()+0.5) for i in range(Ntrials)]
    choice = [int(rand()+0.5) for i in range(Ntrials)]
    color_chosen = [{0: original_ball[i], 1: 0}[choice[i]] for i in range(Ntrials)]
    selected_trials = [i for i in range(Ntrials) if color_chosen[i] == 0]
    remaining_ball = [{1: original_ball[i], 0: 0}[choice[i]] for i in selected_trials]
    probability = 1 - 1.0*sum(remaining_ball)/len(selected_trials)
    
    print '# trials =', Ntrials
    print '# trials when the original ball is white =', Ntrials-sum(original_ball)
    print '# trials when the original ball is selected = ', Ntrials-sum(choice)
    print '# trials when a white ball is selected =', len(selected_trials)
    print '# selected trials when the remaining ball is black =', sum(remaining_ball)
    print 'P(white in the bag | white selected) =', probability


if __name__ == '__main__':
    try:
        main(int(sys.argv[1]))
    except:
        main(100000)