带有列表/标量作为参数的函数

Question

如果我定义一个函数如下：

def somefunc(a, b, c):
    out = [ (a+b)/c for a, b, c in zip(a, b, c) ]
    return out

那么所有的参数必须是长度相同的列表，否则我得到的输出会被截断。举个例子：somefunc([1, 2], [3, 4], [5, 5]) 会给我返回 [0.8, 1.2]。

如果能在某些情况下输入一个单独的数字，而不是列表，那就太方便了，比如：somefunc([1, 2], [3, 4], 5)。当然，任何参数都可以简化为一个单独的值，而不是列表。通过使用numpy，可以用这样的函数定义来实现：

from numpy import array

def somefunc(a, b, c):

    a=array(a)
    b=array(b)
    c=array(c)

    out = (a+b)/c
    return out

有没有办法在不使用numpy的情况下（也不需要一个个转换所有参数）做到这一点呢？

Answer 1

Numpy真是太棒了，非常适合像你提到的那种计算。强烈建议你去用它。

另外，你不需要把所有的参数都转换成数组，只需要一个就可以了。而且最好用asarray，这样可以避免在不需要的时候创建新的副本。

def somefunc(a, b, c):
    return (asarray(a) + b) / c

Answer 2

当然，你可以创建一个无限的列表，这个列表会重复任何不是列表的内容。

from itertools import repeat

def broadcast(o):
    return o if hasattr(o, '__iter__') else repeat(o)

def somefunc(a, b, c):
    out = [ (a+b)/c for a, b, c in zip(broadcast(a), broadcast(b), broadcast(c)) ]
    return out

>>> somefunc([1, 2], [3, 4], 5)
[0.8, 1.2]
>>> somefunc([1, 2], 5, [3, 4])
[2.0, 1.75]
>>>