USACO bisquares集合等差数列问题代码优化求助

你的代码逻辑正确，但存在几处可以大幅提升性能的瓶颈，以下是具体优化方案和改进后的代码：

核心性能瓶颈分析 #

1. bisquares生成时的重复计算：原代码中p和q全范围遍历，重复计算了p²+q²和q²+p²
2. 集合查询的哈希开销：用集合做存在性查询比布尔数组索引慢很多
3. 无效diff的无意义遍历：未利用bisquares的模4性质过滤不可能的公差
4. 堆排序的额外开销：逐个推入堆的效率低于批量排序

具体优化措施 #

1. 优化bisquares生成，减少重复计算 #

仅遍历p ≤ q的组合，避免重复计算对称项，直接减少约50%的计算量：

bisquares = set()
for p in range(pqMax + 1):
    p_sq = p * p
    for q in range(p, pqMax + 1):
        bisquares.add(p_sq + q * q)

2. 用布尔数组替代集合，加速存在性查询 #

布尔数组通过直接索引访问，比集合的哈希查询快数倍：

max_bs = 2 * (pqMax ** 2)
bs_arr = [False] * (max_bs + 1)
for x in bisquares:
    bs_arr[x] = True

3. 利用模4性质过滤无效diff #

bisquares的模4结果只能是0、1、2，因此公差d需满足所有序列项的模4值都在该范围内：

• 若start%4 ==0：d不能是3 mod4
• 若start%4 ==1：d不能是2 mod4
• 若start%4 ==2且length>2：d不能是3 mod4

4. 替换堆排序为批量排序 #

先收集所有有效结果，最后一次性排序，避免堆操作的额外开销。

5. 提前验证终点，减少无效检查 #

在调用完整序列验证前，先检查等差数列的终点是否在bisquares中，快速跳过不可能的diff。

优化后的完整代码 #

with open('ariprog.in', 'r') as fin, open('ariprog.out', 'w') as fout:
    length = int(fin.readline().strip())
    pqMax = int(fin.readline().strip())

    # 生成bisquares集合，避免重复计算
    bisquares = set()
    for p in range(pqMax + 1):
        p_sq = p * p
        for q in range(p, pqMax + 1):
            bisquares.add(p_sq + q * q)
    
    # 转换为有序列表，方便遍历
    bs_list = sorted(bisquares)
    max_bs = bs_list[-1]

    # 创建布尔数组加速存在性查询
    bs_arr = [False] * (max_bs + 1)
    for x in bisquares:
        bs_arr[x] = True

    res = []

    # 处理长度为1的特殊情况
    if length == 1:
        for num in bs_list:
            fout.write(f"{num} 0\n")
        exit()

    # 遍历每个起始点
    for start in bs_list:
        max_diff = (max_bs - start) // (length - 1)
        if max_diff < 1:
            continue
        
        start_mod4 = start % 4
        for diff in range(1, max_diff + 1):
            diff_mod4 = diff % 4
            # 过滤无效公差
            if start_mod4 == 0 and diff_mod4 == 3:
                continue
            if start_mod4 == 1 and diff_mod4 == 2:
                continue
            if start_mod4 == 2 and length > 2 and diff_mod4 == 3:
                continue
            
            # 提前检查终点是否存在
            end = start + diff * (length - 1)
            if not bs_arr[end]:
                continue
            
            # 验证整个序列
            valid = True
            current = start + diff
            for _ in range(length - 2):
                if not bs_arr[current]:
                    valid = False
                    break
                current += diff
            if valid:
                res.append((diff, start))
    
    # 排序后输出
    if res:
        res.sort()
        for diff, start in res:
            fout.write(f"{start} {diff}\n")
    else:
        fout.write("NONE\n")

额外优化方向 #

对于超大m值（如m>200），可以进一步：

• 预计算bisquares的模4分组，直接从对应组中筛选可能的起始点和公差
• 采用滑动窗口算法在有序bisquares列表中快速定位连续等差数列

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