快速搜索大型结构化 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 对此进行了测试。在我的系统上最多花费了几秒钟。