求两个平方数的第 n 个最小和

第一个算法想法是使用堆。

凭直觉，您知道最小的 a 和 b 是最小的，最小的 a²+b² 将是。所以：

• 一目了然，10² + 3² < 10² + 4²
• 一目了然，10²+3²是否<？7²+8²

我建议将 （a，b） 对添加到最小堆中，慢慢增加 a 和 b。然后，序列中的下一个元素是堆中的下一个元素。

在下面的算法中，表示整数除以 2;例如，和 都等于 。/210/211/25

function nth_sum_of_two_nonzero_squares(n):
    heap <- empty min-heap
    for k in range [2 .. n+1]:
        for a in range [1 .. k/2]:
            push a² + (k-a)² to heap if it's not already in the heap
        next_term <- pop(heap)
    return next_term

略有改进：为给定添加的最小项是 。我们可以利用该值的知识来避免向堆添加更多候选者，只要堆的当前根仍然小于下一个候选者。ka = k / 2

这为我们提供了：

function nth_sum_of_two_nonzero_squares(n):
    heap <- empty min-heap
    i = 0
    k = 2
    push 1²+1² to heap
    loop forever:
        next_candidate = ((k+1)/2)² + ((k+2)/2)² = ((k+1)²+1)/2
        while next_candidate > peek(heap):
            i <- i + 1
            next_term = pop(heap)
            if i == n:
                return next_term
        k <- k + 1
        for a in range [1 .. k/2]:
            if a² + (k-a)² not already in heap:
                push a² + (k-a)² to heap
        i <- i+1
        next_term = pop(heap)
        if i == n:
            return next_term

缺点是，我们为弹出的每个术语向堆中添加越来越多的术语。在这里，堆的大小看起来是二次的。但是，我确信 的最大值是 的平方根的量级。如果这是真的，那么算法只使用伪线性时间和线性空间。kkn

您可以通过尝试为弹出的每个术语向堆中添加更少的术语来改进此算法。这需要更多的直觉来了解如何以 a²+b² 的伪递增顺序对 （a， b） 进行排序。

在 python 中测试并与 OEIS 进行比较：正确性为 2 个非零平方之和的数字：

from heapq import heappush, heappop
from itertools import count, islice

def sums_of_two_nonzero_squares():
    heap = [2] # 2 == 1² + 1²
    seen = {2}
    for k in count(3):
        next_candidate = (k*k + 1)//2
        while heap and heap[0] < next_candidate:
            next_term = heappop(heap)
            yield next_term
            seen.remove(next_term)
        for a in range(1, k//2+1):
            x = a**2+(k-a)**2
            if x not in seen:
                seen.add(x)
                heappush(heap, x)

oeis_groundtruth = [    2, 5, 8, 10, 13, 17, 18, 20, 25, 26, 29, 32, 34, 37, 40, 41, 45, 50, 52, 53, 58, 61, 65, 68, 72, 73, 74, 80, 82, 85, 89, 90, 97, 98, 100, 101, 104, 106, 109, 113, 116, 117, 122, 125, 128, 130, 136, 137, 145, 146, 148, 149, 153, 157, 160, 162, 164, 169, 170, 173, 178]

result = list(islice(sums_of_two_nonzero_squares(), len(oeis_groundtruth)))

print(result)
# [2, 5, 8, 10, 13, 17, 18, 20, 25, 26, 29, 32, 34, 37, 40, 41, 45, 50, 52, 53, 58, 61, 65, 68, 72, 73, 74, 80, 82, 85, 89, 90, 97, 98, 100, 101, 104, 106, 109, 113, 116, 117, 122, 125, 128, 130, 136, 137, 145, 146, 148, 149, 153, 157, 160, 162, 164, 169, 170, 173, 178]

print('result == oeis_groundtruth: ', (result == oeis_groundtruth))
# result == oeis_groundtruth:  True

我不知道你怎么称呼蛮力，但这个问题可能可以通过详尽的搜索在时间O（n）内解决。

考虑一个边 2√n 的正方形，原点处有一个角，计算该正方形内的所有值，并每次在位数组中设置一个位。

这需要与 （2√n）² = 4n 成比例的运算。

---

要使这种方法起作用，必须检查 2√n 是否足够大以生成 n 个第一个毕达哥拉斯和。

这是 https://stackoverflow.com/a/74787074/585411 的变体。

如果它作为答案返回，那么它确实有效并且需要内存。正如我对 @yves-daoust 评论的那样，不是，但可以证明是并且可能是类似的东西。所以这比二次要好得多。mO(m)O(sqrt(m) log(m))mO(n)O(n log(n))O(n log(log(n))))

import math

def nth_two_square(n):
    squares = [1, 4]
    high = 1

    while True:
        low = high + 1
        high = squares[-1]+1
        block = [False] * (high - low + 1)
        i = len(squares) - 1
        j = 0
        while squares[j] + squares[j] <= high:
            while high < squares[j] + squares[i]:
                i -= 1
            k = i
            while j <= k and low <= squares[j] + squares[k]:
                block[squares[j] + squares[k] - low] = True
                k = k-1
            j = j + 1

        for i in range(len(block)):
            if block[i]:
                n -= 1
                if n < 1:
                    return low + i

        to_add = int(math.log(squares[-1]))
        for _ in range(to_add):
            squares.append((len(squares)+1)*(len(squares)+1))

print(nth_two_square(1000000))

如果不是检查我们返回哪个，我们将得到一个可以按顺序返回序列的版本。n

def two_squares():
    squares = [1, 4]
    high = 1

    while True:
        low = high + 1
        high = squares[-1]+1
        block = [False] * (high - low + 1)
        i = len(squares) - 1
        j = 0
        while squares[j] + squares[j] <= high:
            while high < squares[j] + squares[i]:
                i -= 1
            k = i
            while j <= k and low <= squares[j] + squares[k]:
                block[squares[j] + squares[k] - low] = True
                k = k-1
            j = j + 1

        for i in range(len(block)):
            if block[i]:
                yield low + i

        to_add = int(math.log(squares[-1]))
        for _ in range(to_add):
            squares.append((len(squares)+1)*(len(squares)+1))

n = 100
for m in two_squares():
    print(m)
    n -= 1
    if n < 1:
        break

---

这背后的一些理论解释了我对其他答案留下的评论。

形式为 with 和 整数的复数集合称为高斯整数。高斯整数的行为很像整数，它们对素数具有独特的因式分解，等等。然而，高斯素数略有不同。a + biab

首先，素数以四组为一组。如果两个数相同，则认为它们是相同的素数，直到 、 或 的因数。1i-1-i

接下来，如果一个数字是两个平方的总和，它就不能是高斯素数。那是因为.相反，如果我们能将一个数字分解为高斯素数，那么我们就可以列出它的所有除数。它是两个平方的总和，当且仅当我们能找到形式的除数时。(a+bi)(a-bi) = a^2+b^2a+bi

所以这意味着这不是高斯素数。同上也不是。但 IS 是高斯素数。1^2+1^2 = 22 = (1+i)(1-i)5 = (1+4i)(1-4i)7

一般结果来自费马的结果。整数中的素数是高斯素数，当且仅当它除以 3 时余数为 1。所以 、 和 都是高斯素数。（进一步的结果是，所有不是高斯素数的素数都可以用一种且只能一种方式写成平方和。71119251317

现在假设除以一些次。然后有这样的.但是，它必须至少同样频繁地分裂。颠倒论证中的符号，我们还发现除法的频率至少与除法的频率一样多。因此，它分裂和相同的次数。7a+bic+di(c+di)*7^k = a+bi(c-di)*7^k = a-bia-bi7a+bia-bia+bia-bi

由此我们发现，它总是被偶数次除以。同样的论点也适用于所有其他也是高斯素数的素数。（这是我开始挥手的地方，但你可以相信这个论点是可以形成的。大约一半的素数与 全等，典型的大整数可以被其中至少一个奇数次整除。因此，对于大数，作为 2 个平方和的数字的密度会下降。7a^2 + b^23 mod 4

但它只会慢慢下降。素数的密度只像 一样下降，其中一半与 一致。因此，密度下降的速度不能比 快。1/log(n)1 mod 41/(2 log(n))

而且，当然，大多数 2 平方的总和都不是质数。所以它下降的速度比这慢。不过，我对此没有很好的估计，只是一个有根据的猜测。