如何迭代一个可以是整数和数组的变量？

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转换为数组或任何可迭代的方式，以便在bar为数组时不会破坏代码。这该怎么做？

Answer 1

numpy.array有一个可选的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。

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

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

你可以试试这个