Numpy: Как добавить массив (4,) и массив (9, 9) к массиву (9, 9, 4)?

Question

import numpy as np
a = np.array([0, 10, 11, 1])
 
# OUT:
# array([ 0, 10, 11,  1])
 
b = np.arange(81).reshape(9, 9)
# OUT:
# array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8],
#       [ 9, 10, 11, 12, 13, 14, 15, 16, 17],
#       [18, 19, 20, 21, 22, 23, 24, 25, 26],
#       [27, 28, 29, 30, 31, 32, 33, 34, 35],
#       [36, 37, 38, 39, 40, 41, 42, 43, 44],
#       [45, 46, 47, 48, 49, 50, 51, 52, 53],
#       [54, 55, 56, 57, 58, 59, 60, 61, 62],
#       [63, 64, 65, 66, 67, 68, 69, 70, 71],
#       [72, 73, 74, 75, 76, 77, 78, 79, 80]])

Я хочу, чтобы каждый элемент b-массива добавлялся к a-массиву и получал результат, как, например, простой l oop:

result = list()

for row in b:
    ls = list()
    
    for i in row:
        c = a + i
        ls.append(c)

    result.append(ls)
np.array(result)

L oop недостаточно эффективен, есть ли способ использовать правило широковещательной передачи Numpy для получения этого результата? (Numpy array)

Answer 1

Типичный "общий" способ сделать это - индексировать с помощью None. Если вы индексируете массив с помощью None, он создает новую ось и транслирует массив по этой оси.

Итак, если у вас есть a.shape = (4,) и b.shape = (9,9), вы можете добавить следующее:

a[:,None,None] + b[None,:,:]

: в индексе для a - единственная ось a, имеющая размер 4. Остальные две оси, None,None, являются новыми осями, которые транслируются для соответствия b.

Индексы :,: в индексе для b - это оси в массиве с формой (9,9). None - это новая ось, широковещательная передача для соответствия a.

Может быть более короткий способ записать это, но я не могу вспомнить, как работают правила неявного широковещания, поэтому я просто сделайте это явным.

Вы также можете использовать numpy.newaxis, если это яснее, но это имеет тот же эффект.