Fast merge elements of two arrays of arrays only if the element is different than zero-CodePudding

So I have two numpy arrays of arrays

a = [[[1, 2, 3, 4], [3, 3, 3, 3], [4, 4, 4, 4]]]
b = [[[0, 0, 4, 0], [0, 0, 0, 0], [0, 1, 0, 1]]]

Both arrays are always of the same size.

The result should be like

c = [[[1, 2, 4, 4], [3, 3, 3, 3], [4, 1, 4, 1]]]

How can I do that in a very fast way in numpy?

CodePudding user response：

Use numpy.where:

import numpy as np

a = np.array([[1, 2, 3, 4], [3, 3, 3, 3], [4, 4, 4, 4]])
b = np.array([[0, 0, 4, 0], [0, 0, 0, 0], [0, 1, 0, 1]])

res = np.where(b == 0, a, b)
print(res)

Output

[[1 2 4 4]
 [3 3 3 3]
 [4 1 4 1]]

CodePudding user response：

For optimal speed use b criterion directly.

Instead of

np.where(b == 0, a, b)
# array([[1, 2, 4, 4],
#        [3, 3, 3, 3],
#        [4, 1, 4, 1]])

timeit(lambda:np.where(b==0,a,b))
# 2.6133874990046024

better do

np.where(b,b,a)
# array([[1, 2, 4, 4],
#        [3, 3, 3, 3],
#        [4, 1, 4, 1]])

timeit(lambda:np.where(b,b,a))
# 1.5850481310044415