元素如果elif函数在python中使用数组

Question

我有一个定义

def myfunc(a, b):
    if a < (b*10):
        result = a*2
    else:
        result = a*(-1)
    return result

现在,当我使用for循环逐个输入my a和b值时,这显然是完美的,但它需要永远(我已经简化了定义)我从经验中知道将值作为数组传递将加速它起来.

那么如何修改此代码以接受数组.我已经使用了any()和all()命令,但我必须使用它们错误,因为我的函数只吐出一个值而不是一个值数组.

我想要的输出的一个例子是:

>>>a = np.array([1,5,50,500])
>>>b = 1
>>>print myfunc(a, b)
array([-1, -5, 100, 1000])

Answer 1

你可以使用np.where:

def myfunc(a, b):
    return np.where(a < b*10, a*2, -a)    

例如,

In [48]: a = np.array([1, 5, 50, 500])

In [49]: b = 1

In [50]: myfunc(a, b)
Out[50]: array([   2,   10,  -50, -500])

请注意,输出与所需输出不同,但与您发布的代码一致.你当然可以通过扭转不平等来获得所需的输出:

def myfunc(a, b):
    return np.where(a > b*10, a*2, -a)

然后

In [52]: myfunc(a, b)
Out[52]: array([  -1,   -5,  100, 1000])