Mefi Bathroom AskMe

from __future__ import division     # Always produce floating point results
import random


gobase = 3000    # prob of someone going to bathroom on given tick => 1/gobase
nr_people = 25
nr_rooms = 3
bticks = 3  # how many ticks people stay in bathroom

def init():
    # Set or reset starting states of global vars
    global room_occ, freq, r0_occ, r1_occ, r1_given_r0

    # Is room n occupied?
    room_occ = [False]*nr_rooms

    # freq[i] => How many times i rooms were occupied
    freq = [0]*(nr_rooms+1)

    # How many times r0 occupied
    r0_occ = 0

    # How many times r2 occupied
    r1_occ = 0

    # How many times r1 occupied if r0 occupied
    r1_given_r0 = 0


# Person is a list of lists, where for each sublist p...
# p[0] = True/False - Is person in bathroom?
# p[1] = nr of bathroom they are in, or None of not in any
# p[2] = clock tick on which person went, or None if not in any

person = []
for i in range(nr_people):
    person.append([False,None,None])
def goes():
    # Randomly decide whether person will go to bathroom
    # Prob of given person going on any given tick is 1/gobase
    return random.uniform(0.0,1.0) < 1/gobase

def countgoes(n):
    # For validating goes() prob dist only
    cg = 0
    for i in range(n):
        if goes():
            cg +=1
    return cg

def ordertried():
    # Randomly pick order in which to try rooms
    return random.sample(range(nr_rooms),nr_rooms)

def gone(p):
    # Has person p gone to the bathroom?
    return p[0]

def timeup(p,tick):
    # Time for person to come back from bathroom?
    return (p[2] + bticks) < tick



def cycle(tick):
    # Run one clock cycle
    for p in person:
        if not gone(p):
            if goes():
                goroom(p,tick)
        else:
            if timeup(p,tick):
                comeback(p)
    examine(tick)

def examine(tick):
    # Examine what is occupied and count frequencies
    global r0_occ, r1_occ, r1_given_r0
    occnr = sum(room_occ)
    freq[occnr] += 1
    if room_occ[0]:
        r0_occ += 1
        if room_occ[1]:
            r1_given_r0 += 1
    if room_occ[1]:
        r1_occ += 1
    #print "Tick %4d: %d rooms occupied" % (tick,occnr)

def comeback(p):
    #print "Person came back from Room %d" % (p[1])
    room_occ[p[1]] = False
    p[0] = False
    p[1] = None
    p[2] = None


def goroom(p,tick):
    # Try rooms in a random order and go to first empty one
    ot = ordertried()
    for r in ot:
        #print "Checking room %d" % r
        if not room_occ[r]:
            room_occ[r] = True
            p[1] = r
            p[0] = True
            p[2] = tick
            #print "Person went to Room %d" % (r,)
            break
    else:
        # All occupied
        pass

def simul(t1,t2):
    # Run a simul for clock ticks t1 to t2 inclusive
    for t in range(t1,t2+1):
        cycle(t)
    sf()

def sr():
    # Show state of people and rooms; used to validate prog only
    print "person"
    for i,p in enumerate(person):
        print i,p
    print
    print "room_occ"
    print room_occ

def sf():
    # Print counts and probs
    print "sum freq = %d" % sum(freq)
    print "freq"
    print freq
    print
    for (i,f) in enumerate(freq):
        print "p(c%d) = %f" % (i,f/sum(freq))
    print
    print "r0_occ = %d, r1_given_r0 = %d, r1_occ = %d" % (r0_occ, r1_given_r0,
                                                          r1_occ)
    print
    print "p(r1_occ) = %f" % (r1_occ / sum(freq))
    print "p(r1_given_r0) = %f" % (r1_given_r0 / r0_occ)
    print "ratio = %f" % ((r1_given_r0 / r0_occ)/(r1_occ / sum(freq)))

def runsim(n=100000):
    # A quick way to run a standard sim
    init()
    simul(1,n)