Generate (i,j) sequence sorted by their sum in Python

user_207352 0
Viewed 120

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.

generator python
2 Answers

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)]
user_5500940

using itertools module will speed up the code:

n = 3
from itertools import product
print(sorted(product(range(n), repeat=2), key=sum))

output :

[(0, 0), (0, 1), (1, 0), (0, 2), (1, 1), (2, 0), (1, 2), (2, 1), (2, 2)]

Here is the test :

from itertools import product
from timeit import timeit

def func1(n):
    return sorted([(i, j) for i in range(n) for j in range(n)], key=sum)

def func2(n):
    return sorted(product(range(n), repeat=2), key=sum)

print(timeit('func1(500)', globals=globals(), number=100))
print(timeit('func2(500)', globals=globals(), number=100))

output:

6.1365569
4.974886799999999

By looking at number=100 and a n=500 I think the difference is notable.

user_139445240
Related