python中从2个列表中获取一定范围内的所有数字组合的高效算法

Question

假设我有两个列表list_1并且list_2

list_1 = [1, 5, 10] list_2 = [3, 4, 15]

我想要获取包含 list_1 和 list_2 中的元素的元组列表，使得元组中的数字之间的差异在某个常数 c 之下。

例如，假设 c 是 2，那么我的元组将是： [(1, 3), (5, 3), (5, 4)]

当然，可以迭代两个列表并检查两个元素之间的差异是否小于 c，但其复杂度为 n^2，我宁愿降低该复杂度。

Answer 1

以下是来自评论的 Marat 想法的实现：

import bisect

def close_pairs(list1,list2,c):
  #assumes that list2 is sorted
  for x in list1:
    i = bisect.bisect_left(list2,x-c)
    j = bisect.bisect_right(list2,x+c)
    yield from ((x,y) for y in list2[i:j])

list_1 = [1, 5, 10]
list_2 = [3, 4, 15]
print(list(close_pairs(list_1,list_2,2)))
#prints [(1, 3), (5, 3), (5, 4)]

为了证明该策略相对于可能被认为是“天真的”方法的潜在改进，让我们timeit。

import timeit

setup_naive = '''
import numpy
list_a = numpy.random.randint(0, 2500, 500).tolist()
list_b = numpy.random.randint(0, 2500, 500).tolist()
c = 2
def close_pairs(list_a, list_b, c):
    yield from ((x,y) for x in list_a for y in list_b if abs(x-y) <= c)
'''

setup_john_coleman = '''
import bisect
import numpy
list_a = numpy.random.randint(0, 2500, 500).tolist()
list_b = numpy.random.randint(0, 2500, 500).tolist()
c = 2
def close_pairs(list_a, list_b, c):
    list_a = sorted(list_a)
    list_b = sorted(list_b)
    for x in list_a:
        i = bisect.bisect_left(list_b,x-c)
        j = bisect.bisect_right(list_b,x+c)
        yield from ((x,y) for y in list_b[i:j])
'''

print(f"john_coleman: {timeit.timeit('list(close_pairs(list_a, list_b, c))', setup=setup_john_coleman, number=1000):.2f}")
print(f"naive: {timeit.timeit('list(close_pairs(list_a, list_b, c))', setup=setup_naive, number=1000):.2f}")

在一台方便的笔记本电脑上，给出的结果如下：

john_coleman: 0.50
naive: 18.35