Python中使用数组的逐元素if elif函数

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

- Rapid

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

- unutbu

---

非常感谢！虽然这个方法可以工作，但是有另一个评论使用了np.all(a < (b*10))与我的现有定义；为什么我应该使用np.where而不是np.all？...编辑-哎呀，看来我把不等式搞反了，幸运的是这没有对我的问题造成困扰 :) - Rapid

np.all 返回一个布尔值，即 True 或 False。它仅测试 a 的每个元素（和相应的 b 元素）是否满足 a < b*10。如果条件在每个元素上都成立，则返回 True，否则返回 False。我不认为这对你在这种情况下有帮助。 - unutbu

太棒了！我不知道你可以使用np.where和另外两个参数，我一定会更经常地使用它。 - KenHBS

Answer 2

使用列表推导式：

myarray = [1, 5, 50, 500]
b = 1
[myfunc(a, b) for a in myarray]

- gefei

Answer 3

你的函数非常简单，可以完全删除：

arr = [1, 5, 50, 500]
arr = [a * 2 if a < b * 10 else -a for a in arr]

- Eric

---

但我预计列表推导式的速度会比numpy数组操作慢得多。 - mgilson

@mgilson：完全正确。我只是没有NumPy的经验。 - Eric