CODEWITHSUNDEEP
pythonnumpy
okvoyce
okvoyce

Reputation: 153

for loop and indexing in Python

I have a 2d array with four columns as follows:

import numpy as np

 testarray=np.array([[0, 662, 2, 12.0637], 
                     [0, 662, 2, 3.197], 
                     [0, 662, 1, 3.0323], 
                     [1, 600, 2, 14.000], 
                     [1, 600, 2, 13.000],
                     [3, 500, 2, 17],
                     [3, 500, 1, 15]])

For each row, if the value of the third column = 1 then I want to keep this row. I also want to keep all rows which matches the value in column 1. For example, row 3 has column 3 = 1. I want to keep this row, but also rows 1 and 2.

Expected output:

test array=[0, 662, 2, 12.0637], 
           [0, 662, 2, 3.197],
           [0, 662, 1, 3.0323],
           [3, 500, 2, 17]
           [3, 500, 1, 15]]

I have managed to extract the values in column one that I need to keep:

i=1
col3=testarray[testarray[:, 2]==i]
col1=np.delete(col3, np.s_[1:],axis=1)  
print(col1)

I just now need to somehow keep all the rows which have a corresponding value in the first column.

I am not sure if this makes sense? I feel like I need to do some sort of for loop or indexing but don't really know where to start. Please help!!!

Upvotes: 1

Views: 53

Answers (1)

StupidWolf
StupidWolf

Reputation: 46908

You identify the first columns where the third column is 1:

testarray[testarray[:,2]==1,0]
array([0., 3.])

Then its a matter of calling the rows whose 1st column contain these entries:

testarray[np.isin(testarray[:,0],testarray[testarray[:,2]==1,0]),:]

array([[  0.    , 662.    ,   2.    ,  12.0637],
       [  0.    , 662.    ,   2.    ,   3.197 ],
       [  0.    , 662.    ,   1.    ,   3.0323],
       [  3.    , 500.    ,   2.    ,  17.    ],
       [  3.    , 500.    ,   1.    ,  15.    ]])

Upvotes: 3

Related Questions