在 Python 中生成按總和排序的 (i,j) 序列

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 需要 125ms：

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

你可以把你的對想象成 (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)]

Answer 2

使用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我認為差異是顯着的。