NumPy добавить вдоль первой оси

Question

Я хотел бы добавить два массива с разными размерами, просто выполнив одинаковое сложение вдоль первой оси.

Невекторное решение:

x = np.array([[[1,2],[3,4],[5,6]],[[7,8],[9,0],[1,2]],[[3,4],[5,6],[7,8]],[[9,0],[1,2],[3,4]]]) #shape (4,3,2)
y = np.array([[1,2],[3,4],[5,6]]) #shape (3,2)

ans = np.empty(x.shape)
for i in range(x.shape[0]):
    ans[i] = x[i] + y

print(ans) #shape (4,3,2)

Как мне сделать эту трансляцию соответствующим образом?

Answer 1

Благодаря трансляции [numpy-doc] вы можете просто использовать:

x + y

Итак, здесь мы вычисляем элемент по индексу i, j, k как:

х _IJK + у _JK

это дает:

>>> x + y
array([[[ 2,  4],
        [ 6,  8],
        [10, 12]],

       [[ 8, 10],
        [12,  4],
        [ 6,  8]],

       [[ 4,  6],
        [ 8, 10],
        [12, 14]],

       [[10,  2],
        [ 4,  6],
        [ 8, 10]]])
>>> (x + y).shape
(4, 3, 2)

Если вы добавляете два массива вместе, так что первый массив имеет, например, три измерения, а два вторых измерения и последние два измерения первого левого массива равны размерам правого массива, массив справа сторона "урезана". Это означает, что он рассматривается как трехмерный массив, где каждый подмассив равен массиву с правой стороны.

Вы также можете «ввести» дополнительные измерения для y в произвольных позициях, как в этот ответ , чтобы «передать» определенное измерение.