生成 (i,j) 序列，按它们在 Python 中的总和排序

Question

我需要编写一个 Python 生成器来生成 0..N 范围内所有可能的数字对。这些对必须按一对的总和排序。 CPU 是否有可能有效地实施？

序列示例：

(0, 0), (0, 1), (1, 0), (0, 2), (1, 1), (2, 0), (1, 2), (2, 1), (2, 2)

实施不佳，N=1000 需要 125 毫秒：

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)

生成器是首选，因为在很多情况下迭代很快就会停止，所以我希望到时候~零时间消耗。

Answer 1

使用itertools模块将加速代码：

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

输出：

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

这是测试：

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

输出：

6.1365569
4.974886799999999

通过查看 number=100 和 n=500，我认为差异很明显。

Answer 2

您可以将您的对想象成 (x, y) 坐标。您想要生成正方形中所有点的坐标。坐标和（total）相等的点在向右向下的对角线上。

我们只需遍历每条对角线上的所有点即可：

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

生成 (i,j) 序列，按它们在 Python 中的总和排序

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

python

generator