Numpy vectorising

Question

How to numpy vectorize the following python code (for loop)? Any help will be much appreciated

arr1 = np.ndarray(shape = (184,184))
arr2 = np.ndarray(shape = (184,184))
arr3 = np.full((184, 184), 0.0, dtype=float)
for i in range(arr1.size[0]):
    for j in range(arr1.size[1]):
        if arr2[i,j] == 0 or arr1[i,j] == 0:
            arr3[i,j]=0
        elif arr2[i,j] == 255 and arr1[i,j] == 255:
            arr3[i, j] = 255

Answer 1

Leverage vectorized operations with masks and boolean-indexing -

mask1 = (arr1==0) | (arr2==0)
mask2 = (arr1==255) & (arr2==255)

arr3[mask1] = 0
arr3[mask2] = 255

If arr3 is already initialized with zeros, we can skip the arr3[mask1] part, as that's assigning zeros anyway and since there's no other conditional statement, we can directly get arr3 using mask2, like so -

arr3 = 255.0*mask2

Sample run for verification -

In [23]: # Setup input
    ...: np.random.seed(0)
    ...: arr1 = (np.random.rand(184,184)>0.5)*255
    ...: arr2 = (np.random.rand(184,184)>0.5)*255

In [24]: # Run original code
    ...: arr3 = np.full((184, 184), 0.0, dtype=float)
    ...: for i in range(arr1.shape[0]):
    ...:     for j in range(arr1.shape[1]):
    ...:         if arr2[i,j] == 0 or arr1[i,j] == 0:
    ...:             arr3[i,j]=0
    ...:         elif arr2[i,j] == 255 and arr1[i,j] == 255:
    ...:             arr3[i, j] = 255

In [25]: # Run proposed code#1
    ...: out = np.full((184, 184), 0.0, dtype=float)
    ...: mask1 = (arr1==0) | (arr2==0)
    ...: mask2 = (arr1==255) & (arr2==255)
    ...: 
    ...: out[mask1] = 0
    ...: out[mask2] = 255

In [26]: np.allclose(arr3, out) #verify code#1
Out[26]: True

In [27]: np.allclose(arr3, 255.0*mask2) #verify code#2
Out[27]: True