How to create a new numpy 2d array with conditions from the existing 2d array

Question

I have 2d array with 3 columns and N rows. In the third column there are only 0 or 1. I need to create 2 numpy arrays. They both contains first 2 columns of the given matrix, but first array has only rows corresponding to 0 from the third column, and second array has only rows to 1.

I've tried but it failed with dimension problems. I haven't used this kind of format before. onlyNormal_Xtest = np.vstack((onlyNormal_Xtest, Xy[Xy[N_train:, 2] == 0]))

Is it possible to do it faster than following?

onlyNormal_Xtest = np.array([])
Xy_test = Xy[N_train:, :]

    for i in range(np.size(Xy_test, 0)):
        if (Xy_test[i, 2] == 0):
            onlyNormal_Xtest = np.append(onlyNormal_Xtest, Xy_test[i, :2])

Actually it still doesn't work due to dimension problems.

Answer 1

Not sure if I understood your question but here is the code i think you were trying to do

a = np.array([[1,2,0],
             [3,4,1],
             [5,6,0],
             [7,8,1]])

Gives

a[a[:,2]==1][:,:2]
array([[3, 4],
       [7, 8]])

a[a[:,2]==0][:,:2]
array([[1, 2],
       [5, 6]])