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-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

当然，我们可以通过列表理解来加强这一点，使整个函数更短一些：

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

列表理解控制流程

2 个答案: