在 python 的二維數組中搜索一維數組

Question

假設我有一個巨大的 2D 數據庫形狀(1.2mil, 6) 。

我想在big_DB中找到一維數組(1, 6)的索引。 實際上，我一次要搜索 64 個這些向量，形狀為(64, 6) 。

這是我的代碼：

for data in range(64): # I have 64 1d arrays to find
    self.idx = np.where((big_DB == arrays[data]).all(axis=1))

這需要 0.043 秒（對於所有 64 個陣列）。 有沒有更快的方法來做到這一點？ 我的項目將調用搜索 function 超過 40,000 次。

編輯）big_DB 是 itertools.product 的結果，行中唯一，浮點數。

Answer 1

我能夠讓它工作的最快方法是使用 Python 的內置dict類型使用O(1)查找。 您需要預處理您的數據庫，這可能最多需要一兩秒鍾，但是查找 go 從我機器上的 >100ms 到 <50us：對於所有 64 次查找，提高了 2000 倍或更好。 由於我使用 100k 元素的數據庫進行了測試，您可能會得到稍差的結果。 您擁有的較大的數據庫可能會導致更多的 hash 沖突。

為了制作查找哈希表，我將big_DB的每一行轉換為字節 object。 這構成了關鍵。 然后值是每個元素的索引，因為這就是您想要進行查找的方式：

dt = f'V{big_DB.shape[1] * big_DB.dtype.itemsize}'
dict_db = dict(zip(map(np.void.item, np.squeeze(big_DB.view(dt))), range(len(big_DB))))

結果查找很簡單

idx = dict_db[x.view(dt).item()]