python - 快速查找集合中的 numpy 向量

Question

我有一个 numpy 数组，例如:

a = np.array([[1,2],
              [3,4],
              [6,4],
              [5,3],
              [3,5]])

我也有一套

b = set((1,2),(6,4),(9,9))

我想找到集合b中存在的向量的索引，这里是

[0, 2]

但是我使用 for 循环来实现这一点，有没有一种方便的方法来完成这项工作，避免 for 循环？我使用的for循环方法:

record = []
for i in range(a.shape[0]):
    if (a[i, 0], a[i, 1]) in b:
        record.append(i)

Answer 1

您可以使用过滤器:

In [8]: a = np.array([[1,2],
              [3,4],
              [6,4],
              [5,3],
              [3,5]])

In [9]: b = {(1,2),(6,4)}

In [10]: filter(lambda x: tuple(a[x]) in b, range(len(a)))
Out[10]: [0, 2]