如何提取numpy结构化数组中满足某些条件的行？

Question

假设我在 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

numpy 有没有办法提取满足某些条件的行。例如：

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

Answer 1

IIUC，这是一种方法numpy

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

返回：

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

如果您只想获取满足条件的索引，请使用numpy.where：

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

警告：

话虽这么说，numpy 可能不是最适合您的目的的工具。要了解原因，请考虑您的数组的a外观：

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

读起来有点困难...

考虑用于pandas组织表格数据：

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