快速找到集合中的Numpy向量 [英] Find numpy vectors in a set quickly

问题描述

我有一个numpy数组，例如:

I have a numpy array, for example:

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

我也有一套

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

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

I want to find the index of vectors that exist in set b, here is

[0, 2]

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

but I use a for loop to implement this, is there a convinient way to do this job avoiding for loop? The for loop method I used:

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

推荐答案

首先，将集合转换为NumPy数组-

First off, convert the set to a NumPy array -

b_arr = np.array(list(b))

然后，基于 this post ，您将有三种方法.让我们使用第二种方法来提高效率-

Then, based on this post, you would have three approaches. Let's use the second approach for efficiency -

dims = np.maximum(a.max(0),b_arr.max(0)) + 1
a1D = np.ravel_multi_index(a.T,dims)
b1D = np.ravel_multi_index(b_arr.T,dims)    
out = np.flatnonzero(np.in1d(a1D,b1D))

样品运行-

In [89]: a
Out[89]: 
array([[1, 2],
       [3, 4],
       [6, 4],
       [5, 3],
       [3, 5]])

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

In [91]: b_arr = np.array(list(b))

In [92]: dims = np.maximum(a.max(0),b_arr.max(0)) + 1
    ...: a1D = np.ravel_multi_index(a.T,dims)
    ...: b1D = np.ravel_multi_index(b_arr.T,dims)    
    ...: out = np.flatnonzero(np.in1d(a1D,b1D))
    ...: 

In [93]: out
Out[93]: array([0, 2])

这篇关于快速找到集合中的Numpy向量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！