列表理解控制流

Question

我试图利用列表推导来重新创建涉及多个elif语句的函数的结果。

我的程序现在是这样的

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)

此程序生成一个包含1、2、3或4的列表，并具有相应的概率。我想知道是否有一种方法可以利用列表理解来执行相同的任务。

假设我有一个向量x = [1,2,3,4]，我的期望结果，和Prob = [0.5,0.75,0.9,1.0]，第I个事件发生的累积概率。如何使用列表推导得到类似的结果?

我想做一些类似

的事情

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

，但这返回x的所有元素，其中num <= test和我只想要第一个。

我希望这是自和感谢任何帮助。

Answer 1

您可以使用next(iterable)来达到很好的效果:next(outcome for outcome, prob in zip(x, Prob) if num <= prob)将计算与for循环体相同的数字:

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

当然，我们可以用一个列表推导使整个函数更短:

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

通常概率总和为1.0，即probs = [0.5,0.25,0.15,0.1]

你可以做一些非常简单的事情

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

p