如何对既可以是整数又可以是数组的变量进行迭代？

Question

我想创建一个同时接受整数和整数数组的多用途函数，比如像delete这样的Numpy函数：

def foo(bar):   # type(foo) can be integer or an array of integers
    for i in bar:
        print(bar)

问题是当bar是一个int时，这显然会引起一个TypeError: 'int' object is not iterable。我找不到如何将bar转换成数组，或者任何可iterable，在bar是数组的情况下不破坏代码。怎么做？你知道吗

Answer 1

def foo(bar):
    if np.isscalar(bar):
        bar = np.array([bar])
    for i in bar:
        print(bar)

你可以试试这个

Answer 2

您可以使用numpy函数atleast_1d和atleast_2d来获得具有最小维数的输入：

>>> atleast_1d(3)
array([3])

>>> atleast_2d(3)
array([[3]])

>>> atleast_1d(randn(1,1))
array([[-0.97860085]])

Answer 3

^{}有一个可选的ndmin参数。将此值设置为1以确保迭代1d数组：

def foo(bar):   # type(bar) can be integer or an array of integers
    for i in np.array(bar, ndmin=1):
        print(i)

如果输入已经是数组，还可以指定copy=False以避免生成副本。你知道吗

注意，我还修改了您的逻辑：您希望在迭代时打印i，而不是bar。你知道吗