寻找第n小的两平方数之和的高效算法（优于O(n²)暴力枚举）

当然存在优于O(n²)暴力枚举的算法，下面介绍两种常用的高效思路：

1. 最小堆（优先队列）+ 去重法 #

这种方法的核心是用堆维护候选的最小元素，避免暴力生成所有可能值再排序的高复杂度，时间复杂度为O(n log n)，空间复杂度O(n)。

具体步骤：

• 初始化：将第一个元素(1²+1², 1, 1)加入最小堆，同时用一个哈希集合记录已经处理过的(a,b)对（避免重复生成相同的平方和，比如(1,2)和(2,1)对应的平方和相同）。
• 循环n次：
  • 取出堆顶的最小元素，这就是第k小的f(k)（k从1到n）。
  • 生成两个候选元素：((a+1)² + b², a+1, b)和(a² + (b+1)², a, b+1)。
  • 检查这两个候选的(a,b)对是否在哈希集合中，若不在则将其加入堆和集合。

这种方法的优势在于每次只处理当前最小的元素，并有序扩展候选集，不会遗漏更小的可能值，同时去重避免了冗余计算。

2. 二分查找+计数法 #

这种方法通过二分查找确定第n小的数值，再验证该数值对应的平方和数量，时间复杂度同样为O(n log n)，空间复杂度更低（O(1)）。

具体步骤：

• 确定二分范围：左边界为2（即1²+1²），右边界可以设为2*n²（足够覆盖第n小的数值）。
• 二分查找：
  • 取中间值mid，计算有多少组(a,b)满足a² + b² ≤ mid。
  • 计数时使用双指针优化：遍历每个正整数a，找到最大的b使得b² ≤ mid - a²，累加符合条件的b的数量（注意a和b都是正整数）。
  • 如果计数≥n，说明第n小的数≤mid，调整右边界；否则调整左边界，直到找到最小的mid满足计数≥n。

为什么你的贪心更新思路不成立？ #

你尝试通过比较(a+1)² + b²和a² + (b+1)²来更新a、b，本质是一种贪心的单路径扩展，但这种方式会遗漏其他路径的更小值。比如f(4)=1²+3²=10，它是从(1,2)扩展而来，而非(2,2)，单路径的贪心无法覆盖所有可能的候选元素，因此会错过正确的下一个最小值。

内容的提问来源于stack exchange，提问作者user6703592