python - 如何使用另一个 NumPy 数组的值过滤一个 NumPy 数组

Question

我有一个这样的 NumPy 数组:

type = array(["A", "B", "A", "B", "B", "A", "A"])

我想像这样使用这个数组来过滤另一个数组:

data = array([5,4,5],
             [0,4,3],
             [2,1,6],
             [6,7,9],
             [0,1,4],
             [7,9,8],
             [1,4,9])

我希望能够按类型过滤数据，这样:

typeA = array([5,4,5],[2,1,6],[7,9,8],[1,4,9])
typeB = array([0,4,3],[6,7,9],[0,1,4])

知道我应该如何实现它吗？由于第一个数组不是 bool 值，我很难找到解决方案

Answer 1

让我们定义你的数组:

>>> from numpy import array
>>> data = array(([5,4,5], [0,4,3], [2,1,6], [6,7,9], [0,1,4], [7,9,8], [1,4,9]))
>>> type_ = array(["A", "B", "A", "B", "B", "A", "A"])

现在，让我们找出您的类型 A:

>>> A = data[type_=="A"]
>>> A
array([[5, 4, 5],
       [2, 1, 6],
       [7, 9, 8],
       [1, 4, 9]])

还有你的类型B:

>>> B = data[type_=="B"]
>>> B
array([[0, 4, 3],
       [6, 7, 9],
       [0, 1, 4]])

这里，type_=="A" 或 type_=="A" 是 bool 数组。我们使用它们从 data 中选择我们想要的元素。

有关使用 bool 或“掩码”数组进行索引的更多信息，请参阅 NumPy docs .

最后，type 是 python 内置的，最好不要覆盖它。因此，我们在这里使用了 type_。