How to change 2d numpy array's element based on condition

Question

Let's say I have three 2D numpy array and a weighted 1D numpy array that I want to weight this to some specific element in the array

import numpy as np
A= np.array([[2,0,2,2],[0,0,1,1],[1,2,1,0]])
B = np.array([[2,1,1,0],[2,2,0,2],[0,0,0,2]])
C = np.array([[0,0,3,2],[1,2,2,1],[2,3,1,5]])
    
Weight = np.array([0. , 0.5, 2. ])
Weight.reshape(-1,1)

1st condition: when we sum these arrays like below

sum = np.sum(A,axis=0) +np.sum(B,axis=0) + np.sum(C,axis=0) 

for anything that is greater than 10, we select the array like below.

#selecting last 2 columns
A_1 = A[:,np.sum(A,axis=0) +np.sum(B,axis=0) + np.sum(C,axis=0)  > 10]
B_1 = B[:,np.sum(A,axis=0) +np.sum(B,axis=0) + np.sum(C,axis=0)  > 10]
C_1 = C[:,np.sum(A,axis=0) +np.sum(B,axis=0) + np.sum(C,axis=0)  > 10]

(array([[2, 2],
        [1, 1],
        [1, 0]]),
 array([[1, 0],
        [0, 2],
        [0, 2]]),
 array([[3, 2],
        [2, 1],
        [1, 5]]))

2nd condition: when A_1 + B_1 - C_1 <0, we would like to modify the element to (A+B)*weight

diff =  A_1 + B_1 - C_1 <0

(array([[False, False],
        [ True, False],
        [False,  True]]),
 array([[ 0,  0],
        [-1,  2],
        [ 0, -3]]))

For where we have True, we would like to change C_1 to

[[3, 2],
[(1+0)*0.5, 1],
[1, (0+2)*2.0]]

(1+0)*0.5 because 0.5 is at 2nd row and (0+2)*2.0 because 2.0 is at 3rd row

hence at the end C becomes..

array([[0, 0, 3, 2],
        [1, 2, (1+0)*0.5, 1],
        [2, 3, 1, (0+2)*2.0]]))

I would like to achieve this resultant C .. without using a loop Any help would be appreciated.

Answer 1

You can use masks:

AB = A+B
m1 = (AB+C).sum(0)>10
m2 = (AB-C) < 0

m = m1&m2

C = C.astype(float)
C[m] = (AB*Weight)[m]

Updated C :

array([[0. , 0. , 3. , 2. ],
       [1. , 2. , 0.5, 1. ],
       [2. , 3. , 1. , 4. ]])