Python 在数值函数中包含逻辑

我想/希望这个问题有一个简单的答案，但恐怕到目前为止，这是在暗示我

我想实现以下功能，包括逻辑：

def my_func(x):
    if x < 5:
        return 2
    else:
        return 2*x        

结果是： y=[2,2,2,2,10,12,14,18]

我想这可以通过函数中的for循环来完成，以将每个元素的结果追加回某种形式的结果数组，但这似乎有点笨拙

---

有更好的方法吗？

如果您知道您的函数将被传递一个numpy数组，您可以使用

numpy

函数编写它，以矢量化的方式执行您想要的操作：

def my_np_func(x):
    return np.where(x<5, 2, 2*x)

但这可能会更慢，因为函数将被多次调用，而不是一次调用。

签出：

[6]中的

：condlist=[my_x<5，my_x>=5]
在[7]中：choicelist=[2，2*my_x]
在[8]中：np.select（condlist，choicelist）
Out[8]：数组（[2,2,2,2,10,12,14,16,18]）

---

@BrenBarn的解决方案很好，但是

np.select

似乎更一般，因为您可以有多个条件和默认值。如果您可以将您的解决方案框架为单个条件，则一定要使用

np.where

，因为它速度更快。

+1正是

np.where

的用途。使用np.vectorize的解决方案并不好，因为它是一个“假”矢量化（实际上不是并行运行的，只是循环）。请注意，如果不传递数组，或者传递数组标量，

，其中

将返回0D数组。这些可能会引起一些意外，但希望不会比数组标量已经引起的意外更糟。

def my_np_func(x):
    return np.where(x<5, 2, 2*x)

np.vectorize(my_func)(x)

In [6]: condlist = [my_x < 5, my_x >= 5]

In [7]: choicelist = [2, 2*my_x]

In [8]: np.select(condlist, choicelist)
Out[8]: array([ 2,  2,  2,  2, 10, 12, 14, 16, 18])