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

Test which Numpy function argument has more than one element

考虑以下函数:

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

    """
    return a, b, c

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

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

您可以使用 locals() 获得所有参数的 dict,然后使用 sizeargmax 找出最大的参数,如下所示:

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]