如何在numpy二维数组中找到与某个列表匹配的所有元素? [英] How to find all elements in a numpy 2-dimensional array that match a certain list?

问题描述

我有一个二维NumPy数组，例如:

I have a 2-dimensional NumPy array, for example:

array([[1, 1, 0, 2, 2],
       [1, 1, 0, 2, 0],
       [0, 0, 0, 0, 0],
       [3, 3, 0, 4, 4],
       [3, 3, 0, 4, 4]])

我想从该数组中获取某个列表中的所有元素，例如(1、3、4).在示例情况下，期望的结果将是:

I would like to get all elements from that array which are in a certain list, for example (1, 3, 4). The desired result in the example case would be:

array([[1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [3, 3, 0, 4, 4],
       [3, 3, 0, 4, 4]])

我知道我可以做到(按照此处的建议 Numpy:在范围内查找元素):

I know that I can just do (as recommended here Numpy: find elements within range):

np.logical_or(
    np.logical_or(cc_labeled == 1, cc_labeled == 3),
    cc_labeled == 4
)

，但这仅在示例情况下有效.实际上，迭代使用for循环和numpy.logical_or真的很慢，因为可能值的列表以数千为单位(而numpy数组的尺寸大约为1000 x 1000).

, but this will be only reasonably effective in the example case. In reality iteratively using for loop and numpy.logical_or turned out to be really slow since the list of possible values is in thousands (and numpy array has approximately the dimension of 1000 x 1000).

推荐答案

您可以使用

You can use np.in1d -

A*np.in1d(A,[1,3,4]).reshape(A.shape)

此外， np.where 可以使用-

Also, np.where could be used -

np.where(np.in1d(A,[1,3,4]).reshape(A.shape),A,0)

您还可以使用 np.searchsorted 通过使用其可选的'side'参数(输入为left和right)来查找此类匹配项，并注意对于匹配项，searchsorted将使用这两个输入输出不同的结果.因此，np.in1d(A,[1,3,4])的等效项为-

You can also use np.searchsorted to find such matches by using its optional 'side' argument with inputs as left and right and noting that for the matches, the searchsorted would output different results with these two inputs. Thus, an equivalent of np.in1d(A,[1,3,4]) would be -

M = np.searchsorted([1,3,4],A.ravel(),'left') != \
    np.searchsorted([1,3,4],A.ravel(),'right')

因此，最终输出将是-

out = A*M.reshape(A.shape)

请注意，如果未对输入搜索列表进行排序，则需要使用可选参数sorter及其np.searchsorted中的argsort索引.

Please note that if the input search list is not sorted, you need to use the optional argumentsorter with its argsort indices in np.searchsorted.

样品运行-

In [321]: A
Out[321]: 
array([[1, 1, 0, 2, 2],
       [1, 1, 0, 2, 0],
       [0, 0, 0, 0, 0],
       [3, 3, 0, 4, 4],
       [3, 3, 0, 4, 4]])

In [322]: A*np.in1d(A,[1,3,4]).reshape(A.shape)
Out[322]: 
array([[1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [3, 3, 0, 4, 4],
       [3, 3, 0, 4, 4]])

In [323]: np.where(np.in1d(A,[1,3,4]).reshape(A.shape),A,0)
Out[323]: 
array([[1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [3, 3, 0, 4, 4],
       [3, 3, 0, 4, 4]])

In [324]: M = np.searchsorted([1,3,4],A.ravel(),'left') != \
     ...:     np.searchsorted([1,3,4],A.ravel(),'right')
     ...: A*M.reshape(A.shape)
     ...: 
Out[324]: 
array([[1, 1, 0, 0, 0],
       [1, 1, 0, 0, 0],
       [0, 0, 0, 0, 0],
       [3, 3, 0, 4, 4],
       [3, 3, 0, 4, 4]])

运行时测试和验证输出-

Runtime tests and verify outputs -

In [309]: # Inputs
     ...: A = np.random.randint(0,1000,(400,500))
     ...: lst = np.sort(np.random.randint(0,1000,(100))).tolist()
     ...: 
     ...: def func1(A,lst):                         
     ...:   return A*np.in1d(A,lst).reshape(A.shape)
     ...: 
     ...: def func2(A,lst):                         
     ...:   return np.where(np.in1d(A,lst).reshape(A.shape),A,0)
     ...: 
     ...: def func3(A,lst):                         
     ...:   mask = np.searchsorted(lst,A.ravel(),'left') != \
     ...:          np.searchsorted(lst,A.ravel(),'right')
     ...:   return A*mask.reshape(A.shape)
     ...: 

In [310]: np.allclose(func1(A,lst),func2(A,lst))
Out[310]: True

In [311]: np.allclose(func1(A,lst),func3(A,lst))
Out[311]: True

In [312]: %timeit func1(A,lst)
10 loops, best of 3: 30.9 ms per loop

In [313]: %timeit func2(A,lst)
10 loops, best of 3: 30.9 ms per loop

In [314]: %timeit func3(A,lst)
10 loops, best of 3: 28.6 ms per loop

这篇关于如何在numpy二维数组中找到与某个列表匹配的所有元素?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！