2

Consider the following function:

def foo(a, b, c):
    """ Toy function

    """
    return a, b, c

Each of these arguments will be of type numpy.array. I need to efficiently determine which of these arguments has more than one element for use further in the function. I'd like to avoid testing each argument with an if statement as the list can be large and performance is important. Assume that only one argument will have more than one element.

How can I determine which of the input arguments has more than one element?

Jason Strimpel
  • 14,670
  • 21
  • 76
  • 106
  • Can you change how your function is called? If so, you can probably get what you need with kwargs: http://stackoverflow.com/questions/3394835/args-and-kwargs – cfrag Aug 18 '15 at 13:40
  • 1
    What's the problem with using e.g. `numpy.array.size`? It shouldn't take longer than 1 ms on the slowest cpus and is mostly independent of array size. – Daniel Lenz Aug 18 '15 at 13:45
  • I'd like to avoid writing a bunch of `if` statements to test each of the args. – Jason Strimpel Aug 18 '15 at 13:48

1 Answers1

0

You can use locals() to get a dict of all the arguments, then use size and argmax to find which is largest, like so:

import numpy as np

a=np.array([1,])
b=np.array([1,])
c=np.array([1,2,3])

def foo(a,b,c):
    args=locals()
    return args.items()[np.array([i[1].size for i in args.items()]).argmax()][1]

biggest = foo(a,b,c)
print biggest
# [1,2,3]
tmdavison
  • 64,360
  • 12
  • 187
  • 165