时间戳

首页 > 解决方案 > 使标量python函数接受列表和数组

python - 使标量python函数接受列表和数组

问题描述

在 python 中，是否有一种简单有效的方法可以使函数f(x: float)同时接受列表和 numpy 数组作为参数（在这种情况下，我想f逐元素应用并以与发送时相同的格式返回结果）？现在，我只需要一维数组。

作为一个例子（我的真实f更复杂），假设我有：

def f(x):
    return math.log(x) if x > 0 else 0.0

然后这个工作，但不是那么优雅 - 也可能不是那么有效，因为递归（我使用它，因为它允许我只有一个函数）：

def f(x):
    if np.isscalar(x):
        return math.log(x) if x > 0 else 0.0
    elif isinstance(x, np.ndarray):
        return np.array([f(i) for i in x], dtype=float)
    else:
        return [f(i) for i in x]

有没有更好的办法？

标签： pythonnumpy-ndarray

解决方案

如果您需要的不仅仅是代码库中的单个函数定义，那么使用装饰器函数将是一个不错的选择。这确实意味着您必须包装您的代码，但我假设您所说的“只有一个函数”的意思是您确实希望 f(x) 中的函数表达式在所有情况下都相同。

使用您现有的代码，装饰器函数将如下所示：

def elementwise_possible(func):
    def wrapper(x):
        if np.isscalar(x):
            return func(x)
        elif isinstance(x, np.ndarray):
            return np.array([func(i) for i in x], dtype=float)
        else:
            return [func(i) for i in x]
    return wrapper

你会这样写你的函数：

@elementwise_possible
def f(x):
    return math.log(x) if x > 0 else 0.0

结果输出变为

In[2]: A = f(2)
In[3]: A
Out[3]: 0.6931471805599453

In[4]: B = f(np.array([2,3,4]))
In[5]: B
Out[5]: array([0.69314718, 1.09861229, 1.38629436])

In[6]:C = f([5,6,7])
In[7]:C
Out[7]: [1.6094379124341003, 1.791759469228055, 1.9459101490553132]

我认为效率应该是一样的。

推荐阅读