在 python 的二维数组中搜索一维数组

Question

Say I have a massive 2D Database shaped (1.2mil, 6) .

I want to find the index of a 1D array (1, 6) in the big_DB . I actually have 64 of these vectors to search for at a time, shaped (64, 6) .

Here's my code:

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

This takes 0.043 sec (for all 64 arrays). Is there a faster method to do this? My project will call the search function over 40,000 times.

Edit) The big_DB is the result of itertools.product, unique in row, float.

Answer 1

The fastest I've been able to get this to work is using O(1) lookup using Python's builtin dict type. You need to pre-process your DB, which may take a second or two at most, but lookups go from >100ms on my machine to <50us: an improvement by 2000x or better for all 64 lookups. You may get slightly worse results because I tested with a 100k-element database. The larger DB you have may cause more hash collisions.

To make the lookup hash-table, I turned each row of big_DB into a bytes object. This makes up the key. Values are then indices of each element, since that's how you want to do the lookup:

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

The resulting lookup is as simple as

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