CODEWITHSUNDEEP
pythonarraysnumpy
JamieL
JamieL

Reputation: 3

How to extract the entire row and columns when condition met in numpy array

I have a numpy array in this form:

array([['A', 2.2, 9.1],
   ['A', 7.1, 2.3],
   ['B', 4.1, 1.1]],dtype=object)

So I would like to query 'A' and then return all rows and columns with (matching) the string 'A'. Anything not matching the condition is ignored. So the output should be:

form = array([['A', 2.2, 9.1],
              ['A', 7.1, 2.3],dtype=object)

I tried using j = form[np.where(form == 'A')]

which gives array(['A', 'A'], dtype=object). This is not what I want.

Could someone please let me know how I can do this?

Upvotes: 0

Views: 823

Answers (4)

Y.Wang
Y.Wang

Reputation: 153

Quick and dirty way:

import numpy as np
form = np.array([row for row in form_original if row[0] == "A"])

Upvotes: 0

U13-Forward
U13-Forward

Reputation: 71580

Try creating a list comprehension then in it we iterate trough the length of test then check if 'A' is in the list witch is the element with the index of i:

import numpy as np
test = np.array([['A', 2.2, 9.1],
   ['A', 7.1, 2.3],
   ['B', 4.1, 1.1]],dtype=object)
print(test[[i for i in range(len(test)) if 'A' in test[i]]])

Output:

[['A' 2.2 9.1]
 ['A' 7.1 2.3]]

Upvotes: 0

Austin
Austin

Reputation: 26039

You can avoid using where like so:

form = np.array([['A', 2.2, 9.1],
             ['A', 7.1, 2.3],
             ['B', 4.1, 1.1]])

print(form[form[:,0] == 'A'])

# [['A' 2.2 9.1] 
#  ['A' 7.1 2.3]]

Upvotes: 1

DavidG
DavidG

Reputation: 25362

You can slice the array when using np.where() so that only this first column is used:

form = np.array([['A', 2.2, 9.1],
   ['A', 7.1, 2.3],
   ['B', 4.1, 1.1]],dtype=object)

j = form[np.where(form[:,0]=='A')]
print (j)
# [['A' 2.2 9.1]
#  ['A' 7.1 2.3]]

Upvotes: 1

Related Questions