测试哪个Numpy函数参数有多个元素

Question

考虑以下功能：

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

    """
    return a, b, c

这些参数中的每一个都是numpy.array类型。我需要有效地确定哪些参数有多个元素可以在函数中进一步使用。我希望避免使用if语句测试每个参数，因为列表可能很大且性能很重要。假设只有一个参数将包含多个元素。

如何确定哪个输入参数有多个元素？

Answer 1

您可以使用locals()获取所有参数的dict，然后使用size和argmax查找最大的参数，如下所示：

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]