Numpy：从具有 1 行的二维数组中获取一维数组

Question

我有一个复杂的函数，returns 一个二维数组，为简单起见，让我们考虑以下内容：

import numpy as np

def get_array():
    a = np.array([[0, 9, 5]])
    return a

是否有 numpy 命令允许我自动检索单行而不是执行以下操作？

def get_array():
    a = np.array([[0, 9, 5]])
    if a.shape[0] == 1:
        return a[0]
    else:
        return a 
    return a

非常感谢！

Answer 1

np.squeeze(a) 将从 a.

中删除任何单位大小的轴

>>> np.squeeze([1, 2, 3])
array([1, 2, 3])

>>> np.squeeze([[1, 2, 3]])
array([1, 2, 3])

>>> np.squeeze([[1, 2, 3], [4, 5, 6]])
array([[1, 2, 3],
       [4, 5, 6]])

虽然这可能不是您想要的，还有：

>>> np.squeeze([[1], [2]])
array([1, 2])

Answer 2

除了@orlp 的回答之外，如果您注重性能，ravel 似乎比挤压更快：

Numpy：从具有 1 行的二维数组中获取一维数组

Numpy: get 1D array from a 2D array with 1 row

numpy

python-3.x

numpy-ndarray