CODEWITHSUNDEEP
pythonfor-looprepeat
Bill TP
Bill TP

Reputation: 1097

For loop, repetitive calculation in Python

from random import uniform

prob = [0.25,0.30,0.45]

def onetrial(prob):
    u=uniform(0,1)
    if 0 < u <= prob[0]:
        return 11
    if prob[0] < u <= prob[0]+prob[1]:
        return 23
    if prob[0]+prob[1] < u <= prob[0]+prob[1]+prob[2]:
        return 39

print onetrial(prob)

I wonder how to reduce the repetitive part in the def using some for-loop techniques. Thanks.

Upvotes: 0

Views: 1021

Answers (3)

jrennie
jrennie

Reputation: 1957

Assuming you call onetrial frequently, calculate the CDF first to make it a bit faster:

from random import uniform

vals = [11, 23, 39]
prob = [0.25, 0.30, 0.45]
cdf = [sum(prob[0:i+1]) for i in xrange(3)]

def onetrial(vals, cdf):
    u = uniform(0, 1)
    for i in range(3):
        if u <= cdf[i]:
            return vals[i]

You could use bisect to make it even faster.

Upvotes: 0

jgritty
jgritty

Reputation: 11935

I like F.J's answer, but I would use a list of tuples, assuming you can easily do so:

from random import uniform

prob = [(0.25, 11), (0.30, 23), (0.45, 39)]

def onetrial(prob):
    u = uniform(0, 1)
    total_prob = 0
    for i in range(3):
        total_prob += prob[i][0]
        if u <= total_prob:
            return prob[i][1]

Upvotes: 1

Andrew Clark
Andrew Clark

Reputation: 208725

The following is equivalent to your current code and it uses a for loop:

from random import uniform

prob = [0.25, 0.30, 0.45]

def onetrial(prob):
    u = uniform(0, 1)
    return_values = [11, 23, 39]
    total_prob = 0
    for i in range(3):
        total_prob += prob[i]
        if u <= total_prob:
            return return_values[i]

I am a little unclear on the relationship between the values you return and the probabilities, it seems like for your code prob will always have exactly 3 elements, so I made that assumption as well.

Upvotes: 1

Related Questions