CODEWITHSUNDEEP
pythonnumpy
gaggio
gaggio

Reputation: 545

Function with lists/scalars as arguments

If I define a function as:

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

all the arguments must be lists of the same length, otherwise I would get the output truncated. Example: somefunc([1, 2], [3, 4], [5, 5]) would give me [0.8, 1.2].

It would be really convenient to be able to input a single number instead of a list when it doesn't change, like: somefunc([1, 2], [3, 4], 5). Of course any of the arguments can be reduced to a single value instead of a list. This can be achieved using numpy with this function definition:

from numpy import array

def somefunc(a, b, c):

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

    out = (a+b)/c
    return out

Is there a way to do the same without numpy (and without converting all the arguments one by one)?

Upvotes: 0

Views: 650

Answers (2)

shx2
shx2

Reputation: 64328

Numpy is awesome, and is perfect for vectorizing computaions like the one you posted. I'd recommend you do use it.

Also, you don't need to convert all the arguments to arrays, only one, and it is better to use asarray, which would avoid creating a new copy when one isn't necessary.

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

Upvotes: 1

Hans Then
Hans Then

Reputation: 11322

Sure, you can make an infinite list repeating any argument that is not a list in itself.

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]
>>>

Upvotes: 2

Related Questions