I need to write a Python generator that produces all possible pairs of numbers in range 0..N. The pairs must be sorted by a pair's sum. Is it possible to implement that CPU-efficiently?
Sequence example:
(0, 0), (0, 1), (1, 0), (0, 2), (1, 1), (2, 0), (1, 2), (2, 1), (2, 2)
A poor implementation, which takes 125ms for N=1000:
N = 1000
# t1 = time.time()
pairs = [(i, j) for i in range(N) for j in range(N)]
pairs2 = list(sorted(pairs, key=sum))
# t2 = time.time()
# print(f'took {t2 - t1} s, n={n}')
# print(pairs2)
A generator is preferred because there are many cases when the iteration will stop rather soon, so I'd expect ~zero time consumption then.
You can imagine your pairs as (x, y) coordinates. You want to generate the coordinates of all points in a square. The points with equal sums of coordinates (total) are on diagonals going downwards to the right.
We just have to go through all the points on each diagonal:
def pairs_by_sum(n):
for total in range(0, n + 1):
for x in range(0, total + 1):
yield (x, total - x)
for total in range(n+1, 2*n + 1):
for x in range(total - n, n+1):
yield(x, total - x)
print(list(pairs_by_sum(2)))
# [(0, 0),
# (0, 1), (1, 0),
# (0, 2), (1, 1), (2, 0),
# (1, 2), (2, 1),
# (2, 2)]
print(list(pairs_by_sum(3)))
# [(0, 0),
# (0, 1), (1, 0),
# (0, 2), (1, 1), (2, 0),
# (0, 3), (1, 2), (2, 1), (3, 0),
# (1, 3), (2, 2), (3, 1),
# (2, 3), (3, 2),
# (3, 3)]