在numpy數組中有效地查找numpy數組中的值

Question

我正在嘗試創建僅包含特定值的numpy數組的副本。 這是我使用的代碼：

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)

......這正是我想要的。

現在，我希望query_val不僅僅是一個值。 答案在這里： Numpy函數多個條件建議使用邏輯和操作，但這是非常低效的，因為你使用==幾次，創建多個中間結果。

在我的情況下，這意味着我沒有足夠的RAM來做到這一點。 有沒有辦法在原生numpy中以最小的空間開銷正確地執行此操作？

Answer 1

這是使用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)

樣品運行 -

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]])