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

Question

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'

please next time make sure you put the code there is actually working so that the person who is trying to answer doesn't spend time cleaning it. Like wrapping strings around with ' ' ... — D_Serg
– D_Serg, Commented Apr 5, 2018 at 19:51
I edited your question to clean up your code so it's runnable. — sacuL
– sacuL, Commented Apr 5, 2018 at 19:57

sacuL · Accepted Answer · 2018-04-05 20:07:08Z

5

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

edited Apr 5, 2018 at 20:07

answered Apr 5, 2018 at 19:45

sacuL

51.6k9 gold badges88 silver badges115 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

NaN Over a year ago

If you really need to see it column-wise... not as fancy, but better than adding a whole package, consider ... a.reshape(a.shape[0], 1) for a quick view

Collectives™ on Stack Overflow

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

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related