如何在 Python 中同时处理浮点数和数组输入？答案

【问题标题】：How to handle both float and array input in Python?如何在 Python 中同时处理浮点数和数组输入？
【发布时间】：2019-09-24 23:14:58
【问题描述】：

我正在尝试编写一个接受浮点数或浮点数数组的函数，并使用相同的代码行处理它们。例如，如果它是一个浮点数，我想返回浮点数本身，如果它是一个数组，我想返回浮点数数组的总和。像这样的

def func(a):
  return np.sum(a)

让func(1.2) 返回 1.2，func(np.array([1.2,1.3,1.4]) 返回 3.9。

【问题讨论】：

【解决方案1】：

确保输入是 NumPy 数组的常用方法是使用np.asarray()：

import numpy as np

def func(a):
  a = np.asarray(a)
  return np.sum(a)

func(1.2)
# 1.2
func([1.2, 3.4])
# 4.6
func(np.array([1.2, 3.4]))
# 4.6

或者，如果您想获取数组的len()，请确保它至少是一维的，请使用np.atleast_1d()：

def func(a):
  a = np.atleast_1d(a)
  return a.shape[0]

func(1.2)
# 1
func([1.2, 3.4])
# 2
func(np.array([1.2, 3.4]))
# 2

【讨论】：

【解决方案2】：

这已经可以了，问题出在哪里？

import numpy as np
def func(a):
  return np.sum(a)
print(func(np.array([1.2,2.3,3.2])))
print(func(1.2))

输出：

6.7
1.2

【讨论】：

【解决方案3】：

您可以使用参数扁平化：

def func(*args):
    # code to handle args
    return sum(args)

现在以下具有相同的行为：

>>> func(3)
3
>>> func(3, 4, 5)
12
>>> func(*[3, 4, 5])
12

【讨论】：

【解决方案4】：

您可以检查输入是否为浮点数，然后在处理总和之前将其放入列表中：

def func(a):
    if isinstance(a, float):
        a = [a]
    return np.sum(a)

【讨论】：