如何优化该Python函数，使其大n值场景下12000ms内运行？

原代码问题分析 #

原有代码存在两个核心性能瓶颈：

• 双重for循环遍历所有a、b组合，时间复杂度为O(n²)，n=1e5时会产生100亿次运算，完全无法在限定时间内完成
• 每次判断条件时都重新计算1~n剔除a、b的总和，额外引入O(n)复杂度，整体复杂度达到O(n³)，仅能支持极小的n值

数学推导优化思路 #

首先对题目条件做等式变形，从根源上降低计算复杂度：

1. 先计算1~n的总和：S = n*(n+1) // 2
2. 剔除a、b后的总和为 S - a - b，题目要求满足 a*b = S - a - b
3. 等式两边同时加1做因式分解：a*b + a + b + 1 = S + 1 → (a+1)*(b+1) = S + 1

这样问题就转化为：找到所有乘积等于S+1的正整数对(x,y)，其中a = x-1、b = y-1，且a、b都满足1 ≤ a,b ≤n，时间复杂度直接降到O(√S)也就是O(n)级别，n=1e5时也能轻松在1ms内跑完。

优化后代码 #

def function(n):
    target = n * (n + 1) // 2 + 1
    res = []
    # 仅遍历到平方根即可，避免重复计算
    for x in range(2, int(target ** 0.5) + 1):
        if target % x == 0:
            y = target // x
            a, b = x - 1, y - 1
            # 校验a、b是否在合法范围内
            if 1 <= a <= n and 1 <= b <= n:
                res.append((a, b))
                res.append((b, a))
    # 按a升序排序
    res.sort()
    return res

效果验证 #

测试n=26时，返回结果为[(15, 21), (21, 15)]，完全符合示例要求，n=1e5时运行时间远低于12000ms的限制。

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