python - 根据字典中是否存在过滤 numpy 数组

Question

我有一个 numpy ndarray 如下:

import numpy as np
x = np.array([[1, 2, 1], [4, 5, 7], [3, 2, 3]])

我有一本字典，其中保存了一些类 ID，如下所示:

k = {1: None, 2: None, 3: None}

现在，该 numpy 数组的最后一列保存类 ID。所以我想做的是根据字典中是否存在类 ID 来过滤 numpy 数组。因此，过滤该输入数组将给出第一行和第三行，因为 7 不在字典中。

所以我得到的类列为:

cls = x[:, -1]

现在，我不知道如何使用它来过滤 x 数组，而无需循环遍历它并创建另一个数组。

Answer 1

这是 numpy.in1d 的一种方式:

keys = list(k.keys())
res = x[np.in1d(x[:, -1], keys)]

print(res)

[[1 2 1]
 [3 2 3]]