列表推导式控制流

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]

但是这将返回所有满足num <= test的x元素，而我只想要第一个。希望我的表达清楚，感谢您的任何帮助。

- eitanlees

---

如果你只想要第一个出现的结果，你可以使用next函数来包装你的列表推导式，例如：next(x[i] for num in u for i, test in enumerate(Prob) if num <= test)。虽然我必须承认我不太确定我是否正确理解了你的意思。 - UnholySheep

代码行数从来都不是衡量软件质量的好指标。事实上，你的顶部代码块运行速度明显更快，可读性也比提出的答案要高得多。 - msw

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

- Mathias Rav

---

非常聪明，但难以阅读，并且比原始帖子中的链接ifs慢大约6倍。我想zip和两个嵌套的for循环会让你付出代价，所以我进行了测试。 - msw

如果速度很重要，使用Joran Beasley基于NumPy的解决方案。如果您想使用迭代器解决问题（这大致是提问者想要的），那么请改用此解决方案。 - Mathias Rav

Answer 2

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

您可以轻松地执行以下操作

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

如果是我，我会采用以下解决方案哈哈

- Joran Beasley