Swap 2 numpy arrays based on condition from different arrays

user_2532034 0
Viewed 67

I have 4 arrays, A,B,C,D. A and B have shape (n,n) and C/D have shape (n,n,m). I am trying to set it up so that when an element of A is greater than B, that array of length m belongs to C. In essence C_new = np.where(A > B, C,D) , D_new = np.where(A < B , D, C). However this gives me a value error (operands could not be broadcast together with shapes)

I am curious if I can use where here instead of just looping through each element?

Edit: example:

A = np.ones((2,2))
B = 2*np.eye(2)
C = np.ones((2,2,3))
D = np.zeros((2,2,3))
# Cnew = np.where(A > B, C,D)-> ValueError: operands could not be broadcast together with shapes (2,2) (2,2,3) (2,2,3)

The Cnew would be zeros in the (0,0) and (1,1) index.

array-broadcasting numpy python
1 Answers

You need to add a new axis at the end of the condition in order for it to broadcast correctly:

C_new = np.where((A > B)[..., np.newaxis], C, D)
D_new = np.where((A < B)[..., np.newaxis], D, C)
user_29378310
Related