python - 如何制作一个可以处理单个输入或输入列表的函数

Question

我将用几个只是平方数的过于简单的函数来说明我的问题。但这是编写我想知道的函数的一般最佳实践。

假设我有一个想要平方的数字，我可以编写以下函数。

def square (x):
    y = x**2
    return y

如果我有一个数字列表并想获得每个元素平方的列表，我可以使用

def square (x_list):
    y_list = []
    for x in x_list:
        y_list.append( x**2 )
    return y_list

但我只想写一个函数来处理这两种情况。它将查看输入是数字还是列表，并采取相应的行动。我可以使用type检测类型来做到这一点，但我想知道最pythonic的方法是什么。

标签： pythonfunction

Answer 1

如果你使用 numpy，这非常简单。对于多个元素，使其成为一个数组。对于单个元素，函数作为它应用。

import numpy as np

def square(a):
    return a**2

a = np.array([2,3,4,5])
print(square(a)) # array([ 4,  9, 16, 25])

b = 9
print(square(b)) # 81