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

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 维(即标量)，将其提升为一维，在数组上运行该函数，如果设置了标志，则在返回之前将其转换回标量。以你为例:

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)

如果您经常这样做，您可以在函数装饰器中添加所有处理。