Python:如何快速搜索集合中的子字符串？

Question

我有一个包含 ~300.000 个元组的集合

In [26]: sa = set(o.node for o in vrts_l2_5) 
In [27]: len(sa)
Out[27]: 289798
In [31]: random.sample(sa, 1)
Out[31]: [('835644', '4696507')]

现在我想根据公共(public)子字符串查找元素，例如前 4 个“数字”(实际上元素是字符串)。这是我的方法:

def lookup_set(x_appr, y_appr):
    return [n for n in sa if n[0].startswith(x_appr) and n[1].startswith(y_appr)]

In [36]: lookup_set('6652','46529')
Out[36]: [('665274', '4652941'), ('665266', '4652956')]

是否有更高效、更快速的方法？

Answer 1

您可以在 O(log(n) + m) 时间内完成，其中 n 是元组的数量，m 是匹配元组的数量，如果您有能力保留元组的两个排序副本。排序本身将花费 O(nlog(n))，即它会渐进地慢然后你的天真方法，但如果你必须做一定数量的查询(更多log(n)，这几乎肯定是相当小的)它会得到返回。

这个想法是，您可以使用二分法找到具有正确的第一个值和正确的第二个值的候选项，然后将这些集合相交。

但是请注意，您需要一种奇怪的比较:您关心以给定参数开头的所有 字符串。这只是意味着在搜索最右边的匹配项时，您应该用 9 填充键。

一个完整的工作(虽然没有经过太多测试)代码:

from random import randint
from operator import itemgetter

first = itemgetter(0)
second = itemgetter(1)

sa = [(str(randint(0, 1000000)), str(randint(0, 1000000))) for _ in range(300000)]
f_sorted = sorted(sa, key=first)
s_sorted = sa
s_sorted.sort(key=second)
max_length = max(len(s) for _,s in sa)

# See: bisect module from stdlib
def bisect_right(seq, element, key):
    lo = 0
    hi = len(seq)
    element = element.ljust(max_length, '9')
    while lo < hi:
        mid = (lo+hi)//2
        if element < key(seq[mid]):
            hi = mid
        else:
            lo = mid + 1
    return lo


def bisect_left(seq, element, key):
    lo = 0
    hi = len(seq)
    while lo < hi:
        mid = (lo+hi)//2
        if key(seq[mid]) < element:
            lo = mid + 1
        else:
            hi = mid
    return lo


def lookup_set(x_appr, y_appr):
    x_left = bisect_left(f_sorted, x_appr, key=first)
    x_right = bisect_right(f_sorted, x_appr, key=first)
    x_candidates = f_sorted[x_left:x_right + 1]
    y_left = bisect_left(s_sorted, y_appr, key=second)
    y_right = bisect_right(s_sorted, y_appr, key=second)
    y_candidates = s_sorted[y_left:y_right + 1]
    return set(x_candidates).intersection(y_candidates)

以及与您的初始解决方案的比较:

In [2]: def lookup_set2(x_appr, y_appr):
   ...:     return [n for n in sa if n[0].startswith(x_appr) and n[1].startswith(y_appr)]

In [3]: lookup_set('123', '124')
Out[3]: set([])

In [4]: lookup_set2('123', '124')
Out[4]: []

In [5]: lookup_set('123', '125')
Out[5]: set([])

In [6]: lookup_set2('123', '125')
Out[6]: []

In [7]: lookup_set('12', '125')
Out[7]: set([('12478', '125908'), ('124625', '125184'), ('125494', '125940')])

In [8]: lookup_set2('12', '125')
Out[8]: [('124625', '125184'), ('12478', '125908'), ('125494', '125940')]

In [9]: %timeit lookup_set('12', '125')
1000 loops, best of 3: 589 us per loop

In [10]: %timeit lookup_set2('12', '125')
10 loops, best of 3: 145 ms per loop

In [11]: %timeit lookup_set('123', '125')
10000 loops, best of 3: 102 us per loop

In [12]: %timeit lookup_set2('123', '125')
10 loops, best of 3: 144 ms per loop

如您所见，此解决方案比您的原始方法快 240-1400 倍(在这些示例中)。

如果你有一大堆比赛:

In [19]: %timeit lookup_set('1', '2')
10 loops, best of 3: 27.1 ms per loop

In [20]: %timeit lookup_set2('1', '2')
10 loops, best of 3: 152 ms per loop

In [21]: len(lookup_set('1', '2'))
Out[21]: 3587
In [23]: %timeit lookup_set('', '2')
10 loops, best of 3: 182 ms per loop

In [24]: %timeit lookup_set2('', '2')
1 loops, best of 3: 212 ms per loop

In [25]: len(lookup_set2('', '2'))
Out[25]: 33053

如您所见，即使匹配项数量约为总大小的 10%，此解决方案的速度也更快。但是，如果您尝试匹配所有 数据:

In [26]: %timeit lookup_set('', '')
1 loops, best of 3: 360 ms per loop

In [27]: %timeit lookup_set2('', '')
1 loops, best of 3: 221 ms per loop

它变得(不是那么多)慢，尽管这是一个非常特殊的情况，我怀疑你会经常匹配几乎所有元素。

请注意，排序数据所花费的时间非常小:

In [13]: from random import randint
    ...: from operator import itemgetter
    ...: 
    ...: first = itemgetter(0)
    ...: second = itemgetter(1)
    ...: 
    ...: sa2 = [(str(randint(0, 1000000)), str(randint(0, 1000000))) for _ in range(300000)]

In [14]: %%timeit
    ...: f_sorted = sorted(sa2, key=first)
    ...: s_sorted = sorted(sa2, key=second)
    ...: max_length = max(len(s) for _,s in sa2)
    ...: 
1 loops, best of 3: 881 ms per loop

如您所见，完成两个排序的副本只需要不到一秒钟的时间。实际上，上面的代码会稍微快一些，因为它对第二个副本进行“就地”排序(尽管 tim-sort 仍然需要 O(n) 内存)。

这意味着如果您必须执行超过 6-8 个查询，此解决方案会更快。

---

注意:python 的标准库提供了一个bisect 模块。但是它不允许 key 参数(尽管我记得读过 Guido 想要它，所以它可能会在将来添加)。因此，如果您想直接使用它，则必须使用“decorate-sort-undecorate”习惯用法。

代替:

f_sorted = sorted(sa, key=first)

你应该这样做:

f_sorted = sorted((first, (first,second)) for first,second in sa)

即您显式插入键作为元组的第一个元素。之后，您可以使用 ('123', '') 作为传递给 bisect_* 函数的元素，它应该会找到正确的索引。

我决定避免这种情况。我复制粘贴模块源中的代码并对其稍作修改，以便为您的用例提供更简单的界面。

---

最后的评论:如果您可以将元组元素转换为整数，那么比较会更快。但是，大部分时间仍然会用于执行集的交集，所以我不知道它到底能提高多少性能。