如何根据参数值选择不同算法编写函数？解决数组ValueError问题

当传入numpy数组作为参数时，x>2会生成与输入同形状的布尔数组，而Python的if语句需要单个布尔值来判断分支，因此会触发ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()错误。

以下是几种可行的解决方法：

方法一：使用np.where（推荐，性能最优） #

np.where可以直接对数组元素进行条件分支处理，属于向量化操作，性能远高于循环：

import numpy as np

def fun(x):
    return np.where(x > 2, np.cos(x), np.sin(x))

该函数会遍历数组的每个元素，满足x>2的元素返回cos(x)，否则返回sin(x)，最终输出同形状的结果数组。

方法二：使用np.vectorize快速适配标量函数 #

如果想保留原有标量逻辑的写法，可以用np.vectorize将标量函数转换为可处理数组的函数：

import numpy as np

def scalar_fun(x):
    if x > 2:
        return np.cos(x)
    else:
        return np.sin(x)

fun = np.vectorize(scalar_fun)

注意：np.vectorize本质是对数组元素做循环遍历，性能不如np.where，适合快速改造原有代码的场景。

方法三：手动遍历数组元素（不推荐，性能差） #

如果需要更直观的逻辑展示，可以手动遍历数组元素并处理：

import numpy as np

def fun(x):
    result = []
    for elem in x:
        result.append(np.cos(elem) if elem > 2 else np.sin(elem))
    return np.array(result)

这种方式代码可读性高，但处理大数据组时性能低下，仅适合小数据量或调试场景。

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