我有一个函数scalar_func(*args)
,它接受可变数量的标量。它对它们进行一些数学运算并输出一个标量。作为一个简单的例子,我们假设 scalar_func 乘以每个数字:
def scalar_func(*args):
out = 1
for arg in args:
out *= arg
return out
我想让 scalar_func
处理列表。为此,我创建了另一个函数 list_func(*args)
。它接受可变数量的列表并创建一个新列表,如下所示:
def list_func(*args):
out = []
for i in range(len(arg[0])):
out.append(scalar_func(arg[0][i], arg[1][i], arg[2][i]...)
return out
显然,这个函数只是伪代码。如何实现list_func
?
最佳答案
您可以使用zip
这里:
def scalar_func(*values):
return sum(values)
def list_func(*args):
out = []
L = list(zip(*args))
for i in range(len(args[0])):
out.append(scalar_func(*L[i]))
return out
list_func([0, 1, 2], [3, 4, 5]) # [3, 5, 7]
如果您有大型列表,您可能希望创建一个迭代器并使用 next
来减少内存消耗:
def list_func(*args):
out = []
L = iter(zip(*args))
for i in range(len(args[0])):
out.append(scalar_func(*next(L)))
return out
这也可以重写以提高效率:
def list_func(*args):
return [scalar_func(*i) for i in zip(*args)]
或者,您可以itertools.starmap
对于等效功能:
from itertools import starmap
def list_func(*args):
return list(starmap(scalar_func, zip(*args)))
关于python - 迭代变量参数的元素,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/52901112/