使用复杂的 if 语句来处理数组？ Python答案

【问题标题】：Get a complicated if-statement to work with an array? Python使用复杂的 if 语句来处理数组？ Python
【发布时间】：2012-12-05 14:51:37
【问题描述】：

我想做这样的事情：

def fun(a,b,c):
    if (a<b**2) & (a<b*c):
        result = a/math.pi
    elif (a<b**2) & (a>=b*c):
        result = b*2/math.pi
    elif (a>=b**2) & (a<b*c):
        result = c*exp(1)
    elif (a>=b**2) & (a>=b*c):
        result = a*b*c*math.pi
    return result,

但是我将如何让它与一个 numpy 数组一起工作呢？数组将是 a、b 和 c 将是单个数字。

我知道 numpy.where，但我只是不明白如何让它像这段代码那样执行。

【问题讨论】：

标签： python arrays if-statement numpy function

【解决方案1】：

你可以嵌套几个np.where，广播应该负责将数组和数字顺利混合：

result = np.where((a < b**2) & (a < b * c), a / np.pi,
                  np.where((a < b**2) & (a >= b * c), b * 2 / np.pi,
                           np.where((a >= b**2) & (a < b*c), c * np.exp(1),
                                    a * b * c * np.pi)))

例如：

>>> a = np.arange(12).reshape(3,4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> b = 1
>>> c = 2
>>> np.where((a < b**2) & (a < b * c), a / np.pi,
             np.where((a < b**2) & (a >= b * c), b * 2 / np.pi,
                      np.where((a >= b**2) & (a < b*c), c * np.exp(1),
                               a * b * c * np.pi)))
array([[  0.        ,   5.43656366,  12.56637061,  18.84955592],
       [ 25.13274123,  31.41592654,  37.69911184,  43.98229715],
       [ 50.26548246,  56.54866776,  62.83185307,  69.11503838]])

【讨论】：

哇，谢谢。我没有意识到我可以这样嵌套它们。这当然让我的生活更轻松。出于兴趣，为什么要使用 np.pi 而不是 math.pi？
我在使用 numpy 时通常不会明确地 import math，但 math.pi 应该完全一样。