How to add a 1d array to a 2d array element-wise to get a 3d array in numpy-CodePudding

I have a 2d array of values, and I want to add a 1d array to this 2d array element wise such that I would get a 3d array where each element is the original 2d array plus a respective element of the 1d array. For example:

A = np.array([
    [10, 9, 8, 7, 6],
    [5, 4, 3, 2, 1]
])
B = np.array([1, 2, 3])

#What A   B should return:
np.array([
    [[11, 10, 9, 8, 7], [6, 5, 4, 3, 2]],
    [[12, 11, 10, 9, 8], [7, 6, 5, 4, 3]],
    [[13, 12, 11, 10, 9], [8, 7, 6, 5, 4]]
])

I was able to do this pretty easily with a normal for loop but is this possible in pure numpy?

CodePudding user response：

I believe this gives you the output you're after?

import numpy as np

A = np.array([
    [10, 9, 8, 7, 6],
    [5, 4, 3, 2, 1]
])
B = np.array([1, 2, 3])

A = A.reshape(1, 2, 5)
B = B.reshape(3, 1, 1)

for each in A   B:
    print (each)
    
# Result:
    # [[11 10  9  8  7]
    #  [ 6  5  4  3  2]]
    # [[12 11 10  9  8]
    #  [ 7  6  5  4  3]]
    # [[13 12 11 10  9]
    #  [ 8  7  6  5  4]]

CodePudding user response：

import numpy as np


A = np.array([
    [10, 9, 8, 7, 6], [5, 4, 3, 2, 1]
])
B = np.array([1, 2, 3])

# What A   B should return:
# np.array([
#     [[11, 10, 9, 8, 7], [6, 5, 4, 3, 2]],
#     [[12, 11, 10, 9, 8], [7, 6, 5, 4, 3]],
#     [[13, 12, 11, 10, 9], [8, 7, 6, 5, 4]]
# ])


temp = np.array([A]*len(B)).flatten()
add = np.repeat(B, len(A.flatten()))

temp  = add

result = temp.reshape((B.shape[0],) A.shape)
print(result)

# np.array([
#     [[11, 10, 9, 8, 7], [6, 5, 4, 3, 2]],
#     [[12, 11, 10, 9, 8], [7, 6, 5, 4, 3]],
#     [[13, 12, 11, 10, 9], [8, 7, 6, 5, 4]]
# ])

CodePudding user response：

you can have fun with list comprehension here and do it with a one-liner

import numpy as np

A = np.array([
    [10, 9, 8, 7, 6],
    [5, 4, 3, 2, 1]
])
B = np.array([1, 2, 3])

r = np.array([A b for b in B])
print(r)

# [[[11 10  9  8  7]
#   [ 6  5  4  3  2]]
# 
#  [[12 11 10  9  8]
#   [ 7  6  5  4  3]]
# 
#  [[13 12 11 10  9]
#   [ 8  7  6  5  4]]]