使用np.linspace时如何实现带条件分支的多值函数？

问题 #

我有一个带条件分支的函数，想结合np.linspace使用，但运行时触发ValueError。原以为对numpy数组应用函数和列表的map操作逻辑一致，示例代码如下：

def f(n):
    if n < 3 or n > 5:
        return 0
    else:
        return n

a = list(range(0, 11))
print(a)
print(list(map(f, a)))

b = np.linspace(0, 10, 11)
print(b)
print(f(b))

列表场景运行正常，但numpy数组场景报错：

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

不想用any()或all()，求正确实现方式。

解决方案 #

1. 向量化函数（np.vectorize） #

将普通函数转换为可处理numpy数组的向量化版本，实现逐元素逻辑应用：

import numpy as np

def f(n):
    if n < 3 or n > 5:
        return 0
    else:
        return n

f_vec = np.vectorize(f)
b = np.linspace(0, 10, 11)
print(f_vec(b))  # 输出: [0. 0. 0. 3. 4. 5. 0. 0. 0. 0. 0.]

2. 原生numpy条件判断（推荐，性能最优） #

利用numpy数组级运算直接处理条件，无需Python循环：

import numpy as np

b = np.linspace(0, 10, 11)
result = np.where((b < 3) | (b > 5), 0, b)
print(result)

这种方式完全依托numpy的底层优化，处理大规模数据时效率远高于逐元素操作。

3. 转列表处理后转回数组 #

如果偏好列表map的写法，可先将numpy数组转为列表，处理后再转回数组：

import numpy as np

def f(n):
    if n < 3 or n > 5:
        return 0
    else:
        return n

b = np.linspace(0, 10, 11)
result = np.array(list(map(f, b)))
print(result)

此方法适合小数据量场景，性能不如原生numpy操作。

内容的提问来源于stack exchange，提问作者DDADDA