CODEWITHSUNDEEP
numpy
RandomFellow
RandomFellow

Reputation: 347

np.where: and condition over multiple columns of a 2d array

I would like to use np.where on a array (arr) in order to retrieve the indices corresponding to the following condition: first column value must be 1, third column value must be two, here is my code so far:

arr = np.array([
    [0,0,0],
    [1,0,2],
    [0,0,0],
    [1,0,2]
    ])


print(np.where(arr[:,0]==1 and arr[:,2]==2))

But it produces this error:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Any idea ?

Thank you

Upvotes: 0

Views: 980

Answers (2)

afterdark
afterdark

Reputation: 11

Maybe this little change would help :

    arr = np.array([
    [0,0,0],
    [1,0,2],
    [0,0,0],
    [1,0,2]
    ])

#print(np.where(arr[:,0]==1 and arr[:,2]==2))
print(np.where(arr[:,0]==1) , np.where(arr[:,2]==0))

Output :

array([1, 3], dtype=int64),) (array([0, 2], dtype=int64)

Upvotes: 0

mozway
mozway

Reputation: 260390

This is an issue of operator precedence, you need to use parentheses:

np.where((arr[:,0]==1)&(arr[:,2]==2))

A more generic method (imagine you have 20 columns to compare) would be to use:

np.where((arr[:,[0,2]]==[1,2]).all(1))

output: (array([1, 3]),)

Upvotes: 3

Related Questions