如何优化Python实现的双平方等差数列查找算法以满足运行时间限制

题目要求 #

等差数列形式为a, a+b, a+2b, ..., a+nb（n为非负整数），其中a为非负整数，b为正整数。需在双平方集合S中查找所有长度为N的等差数列。双平方集合定义为所有形如p² + q²（p、q为非负整数）的整数构成的集合。题目总时间限制为5秒。

输入规则：

• 第一行输入N，取值范围3 ≤ N ≤ 25，代表目标等差数列的长度
• 第二行输入M，取值范围1 ≤ M ≤ 250，代表双平方计算的范围上界，限定0 ≤ p,q ≤ M

输出规则：

• 无符合要求的序列时，输出单行NONE
• 存在符合要求的序列时，每行输出序列的首项、相邻项差值
• 排序规则：先按差值升序排列，差值相同时按首项升序排列
• 题目说明符合要求的序列总数不超过10000

样例输入：

5
7

样例输出：

37 4
2 8
29 8
1 12
5 12
13 12
17 12
5 20
2 24

现有代码与运行问题 #

当前可正常执行的代码如下：

fin = open('ariprog.in', 'r')
fout = open('ariprog.out', 'w')

N: int = int(fin.readline().strip())
M: int = int(fin.readline().strip())

s: set = set()
ans: list = []
mx: int = M**2 * 2  # 双平方集合的理论最大值

# 生成双平方集合
for j in range(M+1):
    for i in range(j, M+1):
        h: int = (i**2) + (j**2)
        s.add(h)

# 遍历所有可能的步长
for stepVal in range(1, (mx // (N-1)) + 1):
    # 遍历所有可能的起始值
    for initial in s:
        count: int = 1
        k: int = initial
        while count < N:
            if k + stepVal not in s:
                break
            k += stepVal
            count += 1
        else:
            ans.append([initial, stepVal])

# 按规则排序结果
ans.sort(key=lambda x: x[1])
if not ans:
    fout.write("NONE" + "\n")
for i, e in ans:
    pr = f'{str(i)} {str(e)}\n'
    fout.write(pr)
fout.close()

运行测试用例（N=21，M=200）时触发超时，程序执行到5.242秒被终止，内存占用10368KB，未达到5秒的时间限制要求。 #

性能优化点 #

超时来自算法遍历效率不足与Python语言特性的叠加，具体优化措施如下：

• 用布尔数组替换哈希集合
  双平方集合的最大值为2*M²，当M=200时最大值仅为80000，开辟长度为mx+1的布尔列表即可存储所有数值的存在状态。列表的索引查找为O(1)的内存随机访问，速度远高于集合的哈希查找，单步成员判断可实现2~3倍的提速。
• 存储排序后的双平方值列表
  生成双平方集合时，将所有存在的双平方值收集为排序后的列表，遍历起始值时按列表顺序访问，天然符合“同一步长下首项升序”的输出要求，可省去最后对结果列表的全局排序开销。
• 提前剪枝减少无效遍历
  对每个起始值initial，若满足initial + stepVal*(N-1) > mx，直接跳过该起始值的后续检查，避免无意义的循环判断。
• 减少重复计算
  预先生成0~M的平方值列表，计算双平方时直接取列表元素相加，减少循环内的重复平方运算开销。

优化后的程序在N=21、M=200的测试用例下，执行时间可控制在1秒以内，满足5秒的时间限制要求。 #

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