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

Question

我正在编写一个函数，希望当输入为单个元素时，函数返回一个元素；当输入为数组时，函数返回一个由输出元素组成的数组，且每个输出元素与输入数组中相应位置的元素相关联。以下是一个示例：

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

这段代码返回了错误：

if a<5:

ValueError: 数组的真值有多个，无法确定。请使用 a.any() 或 a.all()

我期望输出结果为：

0
1
2
3
4

我该如何使它按我所解释的方式工作，因为我相信许多Python函数都是这样工作的呢？谢谢。

- Cupitor

---

错误信息很清楚。你期望 if a<5 做什么？ - Daniel Roseman

这次我真的很想知道为什么我会被踩票了？！！ - Cupitor

Answer 1

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)))

一般来说，我不喜欢在函数中添加打印语句（除非该函数仅用于打印）。如果将输入/输出与计算分离，那么您的函数将更具可重用性。

---

通常情况下，编写一个返回不同类型输出的函数是不好的，例如对于某些输入返回标量，对于其他输入返回数组。如果这样做，那么使用此函数的后续代码都必须检查输出是标量还是数组。或者，代码必须非常谨慎地编写以控制发送到函数的输入类型。如果这样做，代码可能会变得非常复杂或非常有缺陷。

编写简单的函数--要么始终返回数组，要么始终返回标量。

- unutbu

---

正如我在问题中所说，这只是一个虚拟的示例主体。在我的原始代码中没有任何打印输出。你解决问题的方法是针对特定问题的。我正在尝试弄清楚当Python或NumPy函数不知道输入形状时，它们如何处理这些情况。 - Cupitor

通常情况下，NumPy数组是这样编写的，因为它们通常处理矩阵！我的情况几乎相同！ - Cupitor

你能给出一个表现出这种行为的NumPy函数的例子吗？ - unutbu

np.dot(5,3) 返回15，如果你输入(2,[3,3])，它将返回[6,6]。 - Cupitor

我猜测np.dot会对其参数调用np.asarray方法将所有内容转换为数组，然后使用广播机制将数组提升为相同的形状。即使np.dot(5,3)返回一个标量，但如果在代码中这样做，你正在滥用NumPy。 np.dot(5,3)比5*3慢250倍。 - unutbu

那就是我在想的。在这种情况下，一般的程序是什么？是你怀疑的那样还是其他方式？ - Cupitor

Answer 2

如果您希望该函数根据给定的输入是列表还是整数而做出反应，请使用以下代码：

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

通过以上步骤，该函数检查给定的输入是否为数组，如果条件成立，则使用第一个块。接下来的if块检查输入是否为整数。否则将执行else块。

- Aswin Murugesh

---

谢谢。这个可以工作，但我想知道内置的或者例如numpy函数是不是也是这样处理这些情况的方式？ - Cupitor

这是一种检查输入类型未知的方法。 - Aswin Murugesh

Answer 3

当我需要处理这种情况时，通常会从一开始就对输入进行np.asarray，如果它是0维（即标量），则设置一个标志并将其提升为1维，在数组上运行函数，如果设置了标志，则在返回之前将其转换回标量。以你的示例为例：

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)

如果您经常这样做，可以将所有的处理过程封装到一个函数装饰器中。

- Jaime

Answer 4

您可以使用 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，都能正确执行。

- Roland Smith