从以向量为参数的函数返回数组列表答案

【问题标题】：Returning list of arrays from a function having as argument a vector从以向量为参数的函数返回数组列表
【发布时间】：2020-01-17 10:24:59
【问题描述】：

我有一个函数比如：

def f(x):
    A =np.array([[0, 1],[0, -1/x]]);
    return A

如果我使用标量，我将获得：

>>x=1
>>f(x)
array([[ 0.,  1.],
       [ 0., -1.]])

如果我使用数组作为输入，我会得到：

>>x=np.linspace(1,3,3)
>>f(x)
array([[0, 1],
       [0, array([-1.        , -0.5       , -0.33333333])]], dtype=object)

其实我想获取一个数组列表，即： A = [A_1,A_2, ..., A_n]

现在我不太关心它是数组数组还是包含多个数组的列表。我知道我可以在 x 中使用 for 循环来做到这一点。但我认为可能还有另一种方法可以做到这一点，而且可能更有效。

所以我想要的输出是这样的：

>>x=np.linspace(1,3,3)
>>r=f(x)
array([[[0, 1],[0,-1]], 
       [[0, 1],[0,-0.5]],
       [[0, 1],[0,-0.33333]]])
>>r[0]
array([[0, 1],[0,-1]])

或类似的东西

>>x=np.linspace(1,3,3)
>>r=f(x)
[array([[0, 1],[0,-1]]), 
 array([[0, 1],[0,-0.5]]),
 array([[0, 1],[0,-0.33333]])]
>>r[0]
array([[0, 1],[0,-1]])

谢谢

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

在您的函数中，我们可以检查给定参数的类型。在这里，如果x 是np.ndarray 的类型，我们将创建我们想要的嵌套列表，否则我们将像以前一样返回输出。

import numpy as np

def f(x):
    if isinstance(x, np.ndarray):
        v = -1/x
        A = np.array([[[0, 1],[0, i]] for i in v])
    else:
        A = np.array([[0, 1],[0, -1/x]])

    return A

x = np.linspace(1,3,3)
print(f(x))

输出：

[[[ 0.          1.        ]
  [ 0.         -1.        ]]

 [[ 0.          1.        ]
  [ 0.         -0.5       ]]

 [[ 0.          1.        ]
  [ 0.         -0.33333333]]]

【讨论】：

【解决方案2】：

你可以这样做：

import numpy as np

def f(x):
    x = np.array([x]) if type(x)==float or type(x)==int else x
    A = np.stack([np.array([[0, 1],[0, -1/i]]) for i in x]);
    return A

第一行处理x 是int 或float 的情况，因为它不是可迭代的。那么：

f(1)
array([[[ 0.,  1.],
        [ 0., -1.]]])

f(np.linspace(1,3,3))
array([[[ 0.        ,  1.        ],
        [ 0.        , -1.        ]],

       [[ 0.        ,  1.        ],
        [ 0.        , -0.5       ]],

       [[ 0.        ,  1.        ],
        [ 0.        , -0.33333333]]])

【讨论】：