測試哪個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]