python - 快速搜索大型结构化 numpy 数组

Question

我有一个结构化的 numpy 格式数组

dataZero = [(1000, 1045),  # ('gid','lpid')
            (2345, 2500),
                 ...    ]

其中约有 130,000 个条目。我还有另一种格式的结构化数组

dataSnap = [(1002,...,...,...),   # ('gid',...)
            (2400,...,...,...),
            (2490,...,...,...),
                 ...          ]

但这包含 200 万个条目。

对于 dataZero 中的每个条目, i ，我想查找 dataSnap 中的条目满足条件dataZero['gid'][i] <= dataSnap['gid'] <= dataZero['lpid'][i] 。这应该允许返回多个条目。它可以返回满足此条件的条目的索引，也可以返回条目本身 - 只要我在合理的时间〜分钟内执行此操作。

在上面的示例中，这将返回，例如:

[some code] -> [[0],[1,2],...]   # if returning indices of dataSnap

[some code] -> [[(1002,...,...,...)],  # if returning the entries themselves
                [(2400,...,...,...),(2490,...,...,...)],
               ...]

我不知道如何快速做到这一点 - for 循环需要很长时间。可能会转换为一组{}会更好？任何建议/解决方案表示赞赏。

Answer 1

这是一个使用 Python 排序列表(而不是数组)的算法。您可能需要调整它才能与 numpy 一起使用。

def helper(dataZero, dataSnap): # assumes arguments are sorted ascending!
    ds_iter = iter(dataSnap)
    curr_snap = next(ds_iter)
    for lo, hi in dataZero:
        while curr_snap[0] < lo:
            curr_snap = next(ds_iter)
        while lo <= curr_snap[0] <= hi:
            yield curr_snap
            curr_snap = next(ds_iter) 

result = list(helper(dataZero, dataSnap))

我使用包含 200 万个随机生成条目的 dataSnap 和包含来自 dataSnap 的 130k 随机选择范围的 dataZero 对此进行了测试。在我的系统上最多花了几秒钟。