Find numpy vectors in a set quickly
<p>I have a numpy array, for example:</p>

<pre><code>a = np.array([[1,2],
              [3,4],
              [6,4],
              [5,3],
              [3,5]])
</code></pre>

<p>and I also have a set</p>

<pre><code>b = set((1,2),(6,4),(9,9))
</code></pre>

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

<pre><code>[0, 2]
</code></pre>

<p>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:</p>

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

<p>For reference, a straight forward list comprehension (loop) answer:</p>

<pre><code>In [108]: [i for i,v in enumerate(a) if tuple(v) in b]
Out[108]: [0, 2]
</code></pre>

<p>basically the same speed as the <code>filter</code> approach:</p>

<pre><code>In [111]: timeit [i for i,v in enumerate(a) if tuple(v) in b]
10000 loops, best of 3: 24.5 Âµs per loop

In [114]: timeit list(filter(lambda x: tuple(a[x]) in b, range(len(a))))
10000 loops, best of 3: 29.7 Âµs per loop
</code></pre>

<p>But this is a toy example, so timings aren't meaningful.</p>

<p>If <code>a</code> wasn't already an array, these list approaches would be faster than the array ones, due to the overhead of creating arrays.</p>

<p>There are some numpy set operations, but they work with 1d arrays.  We can get around that by converting 2d arrays to 1d structured.</p>

<pre><code>In [117]: a.view('i,i')
Out[117]: 
array([[(1, 2)],
       [(3, 4)],
       [(6, 4)],
       [(5, 3)],
       [(3, 5)]], 
      dtype=[('f0', '&lt;i4'), ('f1', '&lt;i4')])
In [119]: np.array(list(b),'i,i')
Out[119]: 
array([(1, 2), (6, 4), (9, 9)], 
      dtype=[('f0', '&lt;i4'), ('f1', '&lt;i4')])
</code></pre>

<p>There is a version of this using <code>np.void</code>, but it's easier to remember and play with this 'i,i' dtype.</p>

<p>So this works:</p>

<pre><code>In [123]: np.nonzero(np.in1d(a.view('i,i'),np.array(list(b),'i,i')))[0]
Out[123]: array([0, 2], dtype=int32)
</code></pre>

<p>but it is much slower than the iterations:</p>

<pre><code>In [124]: timeit np.nonzero(np.in1d(a.view('i,i'),np.array(list(b),'i,i')))[0]
10000 loops, best of 3: 153 Âµs per loop
</code></pre>

<p>As discussed in other recent <code>union</code> questions, <code>np.in1d</code> uses several strategies.  One is based on broadcasting and <code>where</code>.  The other uses <code>unique</code>, <code>concatenation</code>, <code>sorting</code> and differences.</p>

<p>A broadcasting solution (yes, it's messy) - but faster than <code>in1d</code>.</p>

<pre><code>In [150]: timeit np.nonzero((a[:,:,None,None]==np.array(list(b))[:,:]).any(axis=-1).any(axis=-1).all(axis=-1))[0]
10000 loops, best of 3: 52.2 Âµs per loop
</code></pre>

python numpy set vectorization lookup