Suppose I have two lists list_1 and list_2
list_1 = [1, 5, 10]
list_2 = [3, 4, 15]
I want to get a list of tuples containing elements from both list_1 and list_2 such that the difference between the numbers in a tuple is under a some constant c.
E.g. suppose c is 2 then the tuples I would have would be:
[(1, 3), (5, 3), (5, 4)]
Of course one can iterate over both lists and check that the difference between 2 elements is less than c, but that has a complexity of n^2 and I would rather reduce that complexity.
Here is an implementation of the idea of Marat from the comments:
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)]
To demonstrate the potential improvement of this strategy over what might be thought of as the "naive" approach, let's 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}")
On a handy laptop that gives result like:
john_coleman: 0.50
naive: 18.35
If the lists are sorted as your example suggests, then remove the sorting and then this has runtime complexity O(M+N+P) where M and N are the list sizes and P is the number of close pairs. It keeps an index i so that ys[i] is the smallest y-value not too small, and then walks over ys[i:...] as long as they're not too large, yielding each pair.
def close_pairs(xs, ys, c):
xs = sorted(xs)
ys = sorted(ys) + [float('inf')]
i = 0
for x in xs:
while x - ys[i] > c:
i += 1
j = i
while ys[j] - x <= c:
yield x, ys[j]
j += 1
Benchmark results with lists/ranges 1000 times larger than your example:
904.4 ms close_pairs_naive
4.9 ms close_pairs_John_Coleman
1.8 ms close_pairs_Kelly_Bundy
Benchmark code:
from timeit import timeit
import random
import bisect
from collections import deque
def close_pairs_naive(list_a, list_b, c):
yield from ((x,y) for x in list_a for y in list_b if abs(x-y) <= c)
def close_pairs_John_Coleman(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])
def close_pairs_Kelly_Bundy(xs, ys, c):
xs = sorted(xs)
ys = sorted(ys) + [float('inf')]
i = 0
for x in xs:
while x - ys[i] > c:
i += 1
j = i
while ys[j] - x <= c:
yield x, ys[j]
j += 1
funcs = [
close_pairs_naive,
close_pairs_John_Coleman,
close_pairs_Kelly_Bundy,
]
xs = random.choices(range(15000), k=3000)
ys = random.choices(range(15000), k=3000)
c = 2
args = xs, ys, c
expect = sorted(funcs[0](*args))
for func in funcs:
result = sorted(func(*args))
print(result == expect, func.__name__, len(result))
print()
for _ in range(3):
for func in funcs:
t = timeit(lambda: deque(func(*args), 0), number=1)
print('%6.1f ms ' % (t * 1e3), func.__name__)
print()