如何查找與某個列表匹配的numpy二維數組中的所有元素？

Question

我有一個2維NumPy的陣列，例如：

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）。在本例的情況下期望的結果將是：

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: 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）。

Answer 1

您可以使用np.in1d -

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

此外，np.where可以使用 -

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

您也可以通過使用其可選'side'參數與輸入作爲left，並使用np.searchsorted找到這樣的匹配right並注意到對於匹配，searchsorted會輸出與這兩個輸入不同的結果。因此，np.in1d(A,[1,3,4])等效會 -

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

因此，最後的結果將是 -

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

請注意，如果未排序輸入的搜索列表中，您需要使用可選參數sorter，其argsort指數爲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]]) 

運行測試和驗證輸出 -

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 

Answer 2

使用np.in1d：

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

in1d，正如其名字暗示，操作扁平陣列上，因此您需要在手術後重新塑形。