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

时间:2021-03-25 07:55:08

标签: python arrays dataframe numpy indexing

假设我有一个巨大的 2D 数据库,形状为 (1.2mil, 6)

我想在 (1, 6) 中找到一维数组 big_DB 的索引。我实际上一次要搜索 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 个阵列)。有没有更快的方法来做到这一点?我的项目将调用搜索功能超过 40,000 次。

编辑) big_DB 是 itertools.product 的结果,行中唯一,浮动。

1 个答案:

答案 0 :(得分:0)

我能够以最快的速度使用 Python 的内置 O(1) 类型使用 dict 查找。您需要预处理您的数据库,这可能最多需要一两秒钟,但查找从我机器上的 >100ms 到 <50us:所有 64 次查找都提高了 2000 倍或更好。由于我使用 10 万个元素的数据库进行了测试,因此您可能会得到稍差的结果。您拥有的数据库越大,可能会导致更多的哈希冲突。

为了制作查找哈希表,我将 big_DB 的每一行都变成了一个字节对象。这构成了关键。值是每个元素的索引,因为这是您想要查找的方式:

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