python - 在数值函数中包含逻辑

Question

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

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

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

因此，当将 my_x = np.arange(1,10,1) 形式的数组传递给它时:

y = my_func(my_x)

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

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

请问有更好的方法吗？

Answer 1

如果您知道您的函数将传递一个 numpy 数组，您可以使用 numpy 函数编写它，以矢量化方式执行您想要的操作:

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

您还可以对原始函数进行矢量化并应用它:

np.vectorize(my_func)(x)

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