[SOLVED] Comparing array row-wise so as to generalize argmax function

Comparing array row-wise so as to generalize argmax function

Using Python, consider an array X containing 2d data:

X = np.array([x0,y0], ..., [xn,yn])

and three 1d arrays Y_A, Y_B, Y_C of same length as X containing numbers. Finally consider 3 empty arrays A,B,C. How can I fill these empty arrays A,B,C according to the following pseudo-code?

Pseudo-code:

for each i in range(X):
   if Y_A[i] > Y_B[i] and Y_A[i] > Y_C[i]
   store X[i] to array A
   else if Y_B[i] > Y_A[i] and Y_B[i] > Y_C[i]
   store X[i] to array B
   else store X[i] to array C

My effort which does not work:

for each i in range(len(X)):
    if Y_A[i] > Y_B[i] and Y_A[i] > Y_C[i]:
        A = Y_A
    if Y_B[i] > Y_A[i] and Y_B[i] > Y_C[i]:
        B = Y_B
    else:
        C = Y_C

Solution

Maybe try something like this:

import numpy as np

X = np.random.random((20, 2))
Y_A = np.random.random((20))
Y_B = np.random.random((20))
Y_C = np.random.random((20))

A, B, C = [], [], [] 
for i in range(X.shape[0]):
   if Y_A[i] > Y_B[i] and Y_A[i] > Y_C[i]:
     A.append(X[i])
   elif Y_B[i] > Y_A[i] and Y_B[i] > Y_C[i]:
     B.append(X[i])
   else:
     C.append(X[i])

A = np.array(A)
B = np.array(B)
C = np.array(C)

You can, of course, also create empty numpy arrays and fill them while looping if they have the same length as X.