在numpy数组中有效地查找numpy数组中的值

Question

I am trying to create a copy of my numpy array that contains only certain values. This is the code I was using:

A = np.array([[1,2,3],[4,5,6],[7,8,9]])
query_val = 5
B = (A == query_val) * np.array(query_val, dtype=np.uint16)

... which does exactly what I want.

Now, I'd like query_val to be more than just one value. The answer here: Numpy where function multiple conditions suggests using a logical and operation, but that's very space inefficient because you use == several times, creating multiple intermediate results.

In my case, that means I don't have enough RAM to do it. Is there a way to do this properly in native numpy with minimal space overhead?

Answer 1

Here's one approach using np.searchsorted -

def mask_in(a, b):
    idx = np.searchsorted(b,a)
    idx[idx==b.size] = 0
    return np.where(b[idx]==a, a,0)

Sample run -

In [356]: a
Out[356]: 
array([[5, 1, 4],
       [4, 5, 6],
       [2, 4, 9]])

In [357]: b
Out[357]: array([2, 4, 5])

In [358]: mask_in(a,b)
Out[358]: 
array([[5, 0, 4],
       [4, 5, 0],
       [2, 4, 0]])