CODEWITHSUNDEEP
pythonpython-3.xfunctionargumentsargument-unpacking
user276833
user276833

Reputation: 87

Iterating through elements of variable arguments

I have a function scalar_func(*args) that takes a variable number scalar numbers. It preforms some math with them and outputs a scalar. As a trivial example, we'll assume scalar_func multiplies each number:

def scalar_func(*args):
    out = 1    
    for arg in args:
        out *= arg
    return out

I would like to have scalar_func work with lists. To do so, I made another function list_func(*args). It takes in a variable number of lists and makes a new one like so:

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

Obviously, this function is just pseudocode. How can I implement list_func?

Upvotes: 1

Views: 45

Answers (1)

jpp
jpp

Reputation: 164623

You can use zip here:

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]

If you have large lists, you may wish to create an iterator and use next to reduce memory consumption:

def list_func(*args):
    out = []
    L = iter(zip(*args))
    for i in range(len(args[0])):
         out.append(scalar_func(*next(L)))
    return out

This may also be rewritten for greater efficiency:

def list_func(*args):
    return [scalar_func(*i) for i in zip(*args)]

Alternatively, you can itertools.starmap for the functional equivalent:

from itertools import starmap

def list_func(*args):
    return list(starmap(scalar_func, zip(*args)))

Upvotes: 1

Related Questions