List Comprehension Control Flow

Question

I am trying to utilize list comprehension to recreate the results of a function that involves multiple elif statements.

My program is currently like this

import numpy as np  

def myFunction(x):
    result = []
    for num in x:
        if num <= 0.5:
            result.append(1)
        elif num <= 0.75:
            result.append(2)
        elif num <= 0.9:
            result.append(3)
        else:
            result.append(4)

    return result

u = np.random.uniform(0,1,1000)

myFunction(u)

This program produces a list of 1,2,3 or 4 with the appropriate probabilities. I was wondering if there is a way I could utilize list comprehension to perform the same task.

Let's say I was given a vector x = [1,2,3,4], my desired outcomes, and Prob = [0.5,0.75,0.9,1.0], the cumulative probability that the ith event will occur. How can I use list comprehension get a similar result?

I was trying to do something like

[x[i] for num in u for i, test in enumerate(Prob) if num <= test]

but this returns all of the elements of x where num <= test and I only want the first one.

I hope this makes since and thanks for any help.

Answer 1

You can use next(iterable) to great effect: next(outcome for outcome, prob in zip(x, Prob) if num <= prob) will compute the same number as the body of your for-loop:

def myFunction2(x):
    outcomes = [1, 2, 3, 4]
    probs = [0.5, 0.75, 0.9, 1.0]
    result = []
    for num in x:
        o = next(o for o, p in zip(outcomes, probs) if num <= p)
        result.append(o)
    return result

Of course, we can pepper this up with a list comprehension to make the entire function a bit shorter:

def myFunction3(x):
    outcomes = [1, 2, 3, 4]
    probs = [0.5, 0.75, 0.9, 1.0]
    result = [
        next(o for o, p in zip(outcomes, probs) if num <= p)
        for num in x
    ]
    return result

Answer 2

typically probabilities sum to 1.0 ie probs = [0.5,0.25,0.15,0.1]

you can then do something really easy

numpy.random.choice([1,2,3,4],p=probs)

if it was me this is the solution I would use ;P