python - 使用复杂的 if 语句来处理数组？Python

Question

我想做这样的事情：

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 但我只是不知道如何让它像这段代码一样执行。

Answer 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]])