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)

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

Answer 2

结帐np.select：

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

@ BrenBarn的解决方案非常棒，但np.select似乎更为通用，因为您可以拥有多个条件以及默认值。如果您可以将解决方案构建为单一条件，请务必使用np.where，因为它更快。

在数值函数中包含逻辑

2 个答案: