列表理解控制流

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]

但是 returns 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

列表理解控制流

List Comprehension Control Flow

python

list-comprehension

control-flow