python - 根据列值过滤numpy ndarray(矩阵)

Question

这道题是关于根据某些列值过滤NumPy ndarray。

我有一个相当大的 NumPy ndarray (300000, 50)，我根据某些特定列中的值过滤它。我有 ndtypes，所以我可以按名称访问每一列。

第一列名为 category_code，我需要过滤矩阵以仅返回 category_code 在 ("A", "B", “C”)。

结果需要是另一个 NumPy ndarray，其列仍可通过 dtype 名称访问。

这是我现在做的:

index = numpy.asarray([row['category_code'] in ('A', 'B', 'C') for row in data])
filtered_data = data[index]

列表理解如:

list = [row for row in data if row['category_code'] in ('A', 'B', 'C')]
filtered_data = numpy.asarray(list)

不会工作，因为我原来拥有的 dtypes 不再可用。

是否有更好/更 Pythonic 的方法来实现相同的结果？

可能看起来像:

filtered_data = data.where({'category_code': ('A', 'B','C'})

谢谢!

Answer 1

您可以使用基于NumPy 的库，Pandas ，它有一个更通用的 ndarrays 实现:

>>> # import the library
>>> import pandas as PD

创建一些示例数据作为python字典，其键是列名，其值是作为 python 列表的列值；每列一个键/值对

>>> data = {'category_code': ['D', 'A', 'B', 'C', 'D', 'A', 'C', 'A'], 
            'value':[4, 2, 6, 3, 8, 4, 3, 9]}

>>> # convert to a Pandas 'DataFrame'
>>> D = PD.DataFrame(data)

要仅返回 category_code 为 B 或 C 的行，概念上有两个步骤，但可以在一行中轻松完成:

>>> # step 1: create the index 
>>> idx = (D.category_code== 'B') | (D.category_code == 'C')

>>> # then filter the data against that index:
>>> D.ix[idx]

        category_code  value
   2             B      6
   3             C      3
   6             C      3

请注意 Pandas 与 NumPy 之间的索引区别，NumPy 是构建 Pandas 的库。在 NumPy 中，您只需将索引放在方括号内，用“,”指示要索引的维度，并使用“:”指示您想要其他维度中的所有值(列):

>>>  D[idx,:]

在 Pandas 中，您调用数据框的 ix 方法，并将索引仅放在括号内:

>>> D.loc[idx]