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 的方法是什么。

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