Добавление столбцов в матрицу в NumPy

Question

У меня есть эта функция, которая добавляет в матрицу вектор-столбец единиц.

def add_ones(x):
    return np.hstack((np.ones((x.shape[0], 1)), x))

, когда x - это матрица, она отлично работает.Например:

x = np.array([[1, 2], [3, 4]])
y = np.zeros((5, 5))
add_ones(x)
add_ones(y)

дает

[[1. 1. 2.]
 [1. 3. 4.]]
[[1. 0. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0. 0.]
 [1. 0. 0. 0. 0. 0.]]

, однако, когда я передаю ему один вектор:

z = np.zeros(5)
add_ones(z)

это выдает эту ошибку:

ValueError: all the input arrays must have same number of dimensions

Я бы хотел, чтобы он вернулся

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

что мне делать?

Answer 1

Ваша функция создает двумерный массив единиц, верно?

def add_ones(x):
    return np.hstack((np.ones((x.shape[0], 1)), x))

Это нормально, если x тоже 2d.Но когда x равно 1d, это неправильно.Ваша функция должна обратить внимание на форму x и при необходимости создать 1d массив.

Что-то вроде:

def add_ones(x):
    if x.ndim == 1:
        y = np.ones(1, x.dtype)
    elif x.ndim ==2:
        y = np.ones((x.shape[0],1), x.dtype)
    else: < do something else>
    return np.hstack((y, x))

Answer 2

Вы должны добавить другое измерение к z (сделать его двумерной матрицей 5x1):

add_ones(z[:, np.newaxis])
#array([[ 1.,  0.],
 #      [ 1.,  0.],
#       [ 1.,  0.],
#       [ 1.,  0.],
#       [ 1.,  0.]])