CODEWITHSUNDEEP
pythonnumpy
Endrju
Endrju

Reputation: 343

How to extract rows in numpy structured array that satisfy some condition?

lets say I create structured array in numpy: 

name = ['Tom' , 'Jim', 'Alice', 'Alice', 'Greg']
height = [188, 160, 160, 157, 180]
pet = ['dog', 'cat', 'fish', 'dog', 'cat']

a = np.zeros(len(name), dtype=[('name', 'U30'), ('height', 'i'), ('pet', 'U30')])

a['name'] = name
a['height'] = height
a['pet'] = pet

Is there a way in numpy to extract those rows which satisfy some condition. For example:

'height' == 160 and 'pet' == 'cat'

Upvotes: 2

Views: 659

Answers (1)

sacuL
sacuL

Reputation: 51335

IIUC, Here is a way to do it with numpy

a[(a['height'] == 160) & (a['pet'] == 'cat')]

Which returns:

array([('Jim', 160, 'cat')],
      dtype=[('name', '<U30'), ('height', '<i4'), ('pet', '<U30')])

If you want to get just the index at which the conditions are satisfied, use numpy.where:

np.where((a['height'] == 160) & (a['pet'] == 'cat'))
# (array([1]),)

Caveat:

That being said, numpy might not be the best tool for your purposes. To see why, consider what your array a looks like:

>>> a
array([('Tom', 188, 'dog'), ('Jim', 160, 'cat'), ('Alice', 160, 'fish'),
       ('Alice', 157, 'dog'), ('Greg', 180, 'cat')],
      dtype=[('name', '<U30'), ('height', '<i4'), ('pet', '<U30')])

It's kind of hard to read...

Think about using pandas for organizing tabular data:

import pandas as pd
df = pd.DataFrame({'name':name, 'height':height, 'pet':pet})
>>> df

   height   name   pet
0     188    Tom   dog
1     160    Jim   cat
2     160  Alice  fish
3     157  Alice   dog
4     180   Greg   cat

>>> df.loc[(df.height==160) & (df['pet'] == 'cat')]
   height name  pet
1     160  Jim  cat

Upvotes: 5

Related Questions