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

时间:2014-09-16 14:15:19

标签: python arrays search numpy set

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

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

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

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

但这包含200万个条目。

对于dataZeroi中的每个条目,我想在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循环需要很长时间。某些转化为集合{}可能会更好吗?任何建议/解决方案都表示赞赏。

1 个答案:

答案 0 :(得分:2)

这是一个使用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))

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