如何编写一个函数来处理数组作为元素的输入元素并返回数组

Question

我想写一个函数，并且我希望它在输入是元素时返回一个元素，如果输入是数组，则返回一个输出数组输出数组的每个元素都与输入数组中的相同位置相关联。我给一个虚拟的例子：

import numpy as np 
def f(a): 
    if a<5: 
     print a; 
f(np.arange(11)) 

此代码返回错误： 如果< 5：

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

我期望的输出是：

0 
1 
2 
3 
4 

我怎么能让它按照我解释的方式工作，因为我相信很多python函数都以这种方式工作？

谢谢。

Answer 1

当我不得不处理这种事情时，我通常首先在输入的开始处做一个np.asarray，如果它是0维（即标量），则设置一个标志， ，将它提升为1-D，在数组上运行该函数，并在返回之前将其转换回标量，如果该标志已设置。随着您例如：

def f(a): 
    a = np.asarray(a) 
    is_scalar = False if a.ndim > 0 else True 
    a.shape = (1,)*(1-a.ndim) + a.shape 
    less_than_5 = a[a < 5] 
    return (less_than_5 if not is_scalar else 
      (less_than_5[0] if less_than_5 else None)) 

>>> f(4) 
4 
>>> f(5) 
>>> f([3,4,5,6]) 
array([3, 4]) 
>>> f([5,6,7]) 
array([], dtype=int32) 

如果你这样做很多时候，你可以在一个函数装饰添加所有的处理。

Answer 2

import numpy as np 
def f(a): 
    result = a[a<5] 
    return result 

def report(arr): 
    for elt in arr: 
     print(elt) 

report(f(np.arange(11))) 

一般而言，我不喜欢把函数打印报表（除非该函数什么都不做，但打印）。如果你保持I/O从计算分开，那么你的功能会更加重用。

---

通常也未尝编写返回不同类型的输出，诸如用于某些输入标量和用于其它输入的数组的函数。如果你这样做，那么使用这个函数的后续代码将不得不检查输出是标量还是数组。或者，必须非常小心地编写代码来控制发送到该函数的输入类型。如果你这样做的话，代码可能会变得非常复杂或者很麻烦。

写简单的函数 - 总是返回一个数组，或总是返回一个标量的函数。

Answer 3

您可以使用isinstance来检查参数的类型，然后让您的函数采取正确的操作;

In [15]: a = np.arange(11) 

In [16]: isinstance(a, np.ndarray) 
Out[16]: True 

In [17]: b = 12.7 

In [18]: isinstance(b, float) 
Out[18]: True 

In [19]: c = 3 

In [20]: isinstance(c, int) 
Out[20]: True 

In [21]: d = '43.1' 

In [23]: isinstance(d, str) 
Out[23]: True 

In [24]: float(d) 
Out[24]: 43.1 

In [25]: float('a3') 
--------------------------------------------------------------------------- 
ValueError        Traceback (most recent call last) 
<ipython-input-25-caad719e0e75> in <module>() 
----> 1 float('a3') 

ValueError: could not convert string to float: a3 

这种方式，你可以创建一个函数，做正确的事羯羊它被赋予一个str，一个float，一个int，一个list或numpy.ndarray。

Answer 4

如果你想要的功能，以根据给定的输入是否是一个列表，或只是一个int反应，使用：

def f(a): 
    if type(a)==type([]): 
     #do stuff 
    elif type(a)==type(5): 
     #do stuff 
    else 
     print "Enter an int or list" 

通过上述步骤，功能检查，如果给定的输入是一个数组，如果条件成立，则使用第一个块。接下来如果块检查输入是否为int。否则else块被执行