usaco-1.4.3 等差数列

形式化题目

给定正整数 $N$ 与非负整数 $M$ ，定义双平方数集合 $S$ 为所有满足 $x = p^2 + q^2$ ( $0 \leqslant p, q \leqslant M$ ) 的整数组成的集合。

求所有由集合 $S$ 中元素构成的、长度为 $N$ 的严格递增等差数列 $(a, a+b, a+2b, \dots, a+(N-1)b)$ ，其中首项 $a \geqslant 0$ ，公差 $b > 0$ 且均为整数。

输出所有满足条件的数列对应的数对 $(a, b)$ 。排序规则为：优先按公差 $b$ 升序排列；若公差 $b$ 相同，则按首项 $a$ 升序排列。若不存在任何满足条件的等差数列，则输出 NONE。

正解

思路

1. 双平方数集合预处理

双平方数的值域上限为 $2M^2$ 。当 $M \leqslant 250$ 时，最大数值为 $2 \times 250^2 = 125000$ 。我们可以通过枚举 $0 \leqslant p \leqslant q \leqslant M$ ，计算所有的 $p^2 + q^2$ ，并在一个大小为 $2M^2 + 1$ 的布尔查找表（如 Python 中的 bytearray）中标记其存在性。随后将所有为真的数值提取为一个升序排列的列表。在 $M=250$ 时，不同的双平方数个数至多为 $21047$ 个。

2. 枚举策略选择与公差确定

一个等差数列完全由首项 $a$ 和公差 $b$ 确定。数列中所有项都必须属于集合 $S$ ：

若直接枚举公差 $b$ 和首项 $a$ ，由于大部分 $(a, b)$ 的前两项甚至都不在集合中，会产生大量无效判断。
因而更高效的策略是：直接在双平方数列表中枚举前两项。
- 外层循环枚举数列首项 $a$ （位于列表第 $i$ 项）；
- 内层循环枚举数列第二项 $c = a + b$ （位于列表第 $j$ 项， $j > i$ ）；
- 公差天然确定为 $b = c - a > 0$ 。

3. 剪枝策略

为了避免枚举过多元组，可进行以下几层剪枝：

公差上界剪枝：整个数列最大项为 $a + (N-1)b$ 。由于该项不能超过值域上界 $2M^2$ ，公差必须满足：
$b \leqslant \left\lfloor \frac{2M^2 - a}{N - 1} \right\rfloor$
由于列表是单调递增的，一旦枚举到的第二项导致 $b$ 超出该上限，内层循环可直接 break 终止。
末项前置剪枝：等差数列中间项能否全部合法，末项 $a + (N-1)b$ 是最严苛的关卡。先利用布尔查找表在 $O(1)$ 时间内检查末项是否在集合中；若末项不在集合中，直接跳过当前公差，避免逐项遍历中间项。
中间项倒序检验：当末项合法后，从 $k = N-2$ 向下倒序检查各项 $a + kb$ 是否在集合中，只要发现一项不符即立刻停止。

收集全部合法数对后，按 $(b, a)$ 进行排序输出即可。

代码

Python 版：

python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70#!/usr/bin/env python3
# Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
# rbook: -> https://rbook.roj.ac.cn  https://rbook2.roj.ac.cn
# rainboy的学习导航网站: https://idx.roj.ac.cn
# create_at: 2026-04-18 10:00
# update_at: 2026-04-18 10:00

import sys


def build_bisquare_table(max_m: int) -> tuple[bytearray, list[int]]:
    """生成双平方数布尔查找表与有序升序列表。"""
    max_val = 2 * max_m * max_m
    is_bisquare = bytearray(max_val + 1)
    for p in range(max_m + 1):
        p2 = p * p
        for q in range(p, max_m + 1):
            is_bisquare[p2 + q * q] = 1
    # 列表推导提取有序双平方数
    bisquares = [val for val, exists in enumerate(is_bisquare) if exists]
    return is_bisquare, bisquares


def find_arithmetic_progressions(
    n: int, max_m: int, is_bisquare: bytearray, bisquares: list[int]
) -> list[tuple[int, int]]:
    """搜索所有长度为 n 的双平方数等差数列 (a, b)。"""
    max_val = 2 * max_m * max_m
    results: list[tuple[int, int]] = []
    total = len(bisquares)

    for i in range(total):
        first = bisquares[i]
        diff_limit = (max_val - first) // (n - 1)
        for j in range(i + 1, total):
            diff = bisquares[j] - first
            if diff > diff_limit:
                break
            # 先剪枝末项，再倒序检查中间项
            last = first + (n - 1) * diff
            if not is_bisquare[last]:
                continue
            is_valid = True
            for step in range(n - 2, 1, -1):
                if not is_bisquare[first + step * diff]:
                    is_valid = False
                    break
            if is_valid:
                results.append((diff, first))

    results.sort()
    return [(first, diff) for diff, first in results]


def solve() -> None:
    data = iter(map(int, sys.stdin.buffer.read().split()))
    seq_len = next(data)
    max_coord = next(data)

    is_bisquare, bisquares = build_bisquare_table(max_coord)
    ans = find_arithmetic_progressions(seq_len, max_coord, is_bisquare, bisquares)

    if not ans:
        print("NONE")
    else:
        print("\n".join(f"{first} {diff}" for first, diff in ans))


if __name__ == "__main__":
    solve()

C++ 版（同一算法）：

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95/**
 * Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
 * rbook: -> https://rbook.roj.ac.cn  https://rbook2.roj.ac.cn
 * rainboy的学习导航网站: https://idx.roj.ac.cn
 * create_at: 2026-10-04 13:45
 * update_at: 2026-10-04 13:45
 */
#include <cstdio>
#include <vector>
#include <algorithm>

typedef long long ll;

const int MAXV = 125005; // 双平方数上界 2 * 250^2 = 125000

int n, m;
int maxv;                // 本题的双平方数上界 2*M*M
bool in_s[MAXV];         // in_s[x] = x 是否是双平方数（p^2 + q^2）
int s_list[MAXV];        // s_list[] 存所有双平方数（升序），下标从 0 开始
int cnt_s = 0;           // 双平方数个数

struct Ans {
    ll a;
    ll b;
};
std::vector<Ans> ans;    // 所有合法数对 (a, b)

// 预处理双平方数集合：枚举 0 <= p <= q <= M 标记所有 p^2 + q^2
void build_set() {
    for (int p = 0; p <= m; ++p) {
        for (int q = p; q <= m; ++q) {
            int v = p * p + q * q;
            in_s[v] = true;
        }
    }
    // 把双平方数按升序提取进列表
    for (int x = 0; x <= maxv; ++x) {
        if (in_s[x]) {
            s_list[cnt_s++] = x;
        }
    }
}

// 检查首项 a（列表下标 i）、第二项 c（列表下标 j）出发的数列是否合法
bool check(ll a, ll b) {
    // 末项前置剪枝：末项 a + (n-1)*b 是最严苛的关卡，先 O(1) 检查
    ll last = a + (ll)(n - 1) * b;
    if (!in_s[last]) return false;
    // 中间项倒序检验：从 k = n-2 向下检查，发现不符立刻停止
    for (int k = n - 2; k >= 1; --k) {
        if (!in_s[a + (ll)k * b]) return false;
    }
    return true;
}

// 按 (b, a) 升序比较
bool cmp_ans(const Ans &x, const Ans &y) {
    if (x.b != y.b) return x.b < y.b;
    return x.a < y.a;
}

int main() {
    scanf("%d %d", &n, &m);
    maxv = 2 * m * m;

    build_set();

    // 在双平方数列表里枚举前两项，公差 b = c - a 天然确定
    for (int i = 0; i < cnt_s; ++i) {
        ll a = s_list[i];
        // 公差上界：末项 a + (n-1)*b <= maxv
        ll b_limit = (maxv - a) / (n - 1);
        for (int j = i + 1; j < cnt_s; ++j) {
            ll b = (ll)s_list[j] - a; // 公差 = 第二项 - 首项，恒为正
            if (b > b_limit) break;   // 列表递增，b 超上界后面只会更大，直接 break
            if (check(a, b)) {
                Ans t;
                t.a = a;
                t.b = b;
                ans.push_back(t);
            }
        }
    }

    if (ans.empty()) {
        printf("NONE\n");
        return 0;
    }

    sort(ans.begin(), ans.end(), cmp_ans);
    for (int i = 0; i < (int)ans.size(); ++i) {
        printf("%lld %lld\n", ans[i].a, ans[i].b);
    }
    return 0;
}

复杂度

时间复杂度：
- 预处理双平方数集合需枚举 $0 \leqslant p \leqslant q \leqslant M$ ，时间复杂度为 $O(M^2)$ ，当 $M=250$ 时循环约 $3.1 \times 10^4$ 次。
- 搜索等差数列时，虽然双平方数个数 $|S| \leqslant 21047$ ，但严格的公差上限约束 $b \leqslant \frac{2M^2 - a}{N-1}$ 配合末项 $O(1)$ 前置剪枝，过滤掉了绝大多数无效状态。整个搜索过程在 $1$ 秒内即可完成。
- 结果排序至多 $10000$ 个元素，耗时 $O(K \log K)$ ，可忽略不计。
空间复杂度：
- 布尔数组占用 $2M^2 + 1 \leqslant 125001$ 字节（约 $122\text{ KB}$ ）。
- 双平方数列表与答案列表占用空间均为 $O(|S|)$ 级别，峰值内存不足 $16\text{ MB}$ ，远低于 $128\text{ MB}$ 限制。

总结

本题的关键在于利用双平方数的值域范围较小（ $\leqslant 125000$ ）这一特征，通过布尔查找表将集合归属判定优化至 $O(1)$ 。同时，放弃盲目枚举公差，转为直接枚举集合内的前两项作为候选，配合公差值域边界与末项前置检查，成功将原本庞大的搜索树大幅裁剪。

图示解析

以样例 $N=5, M=7$ 为例，双平方数上界为 $2 \times 7^2 = 98$ 。

状态判定流程

flowchart TD
    A["枚举首项 a in S"] --> B["枚举第二项 c in S (c > a)"]
    B --> C["计算公差 b = c - a"]
    C --> D{"b <= (2M^2 - a) / (N - 1) ?"}
    D -- 否 --> E["直接 break 跳出内层循环"]
    D -- 是 --> F{"末项 a + (N-1)b 是否在集合中？"}
    F -- 否 --> G["continue 跳过当前分支（末项剪枝）"]
    F -- 是 --> H["倒序逐项检查 a + kb 是否全在集合中"]
    H --> I{"全部项均在集合中？"}
    I -- 是 --> J["记录合法解 (b, a)"]
    I -- 否 --> K["舍弃"]

样例搜索过程推演（部分）

对于数列长度 $N=5$ ，找到的第一个合法等差数列为 $a=1, b=4$ ：

项数 $k$	表达式 $a + k \cdot b$	数值	对应的 $p^2 + q^2$ 分解	是否在集合中
第 1 项 ( $k=0$ )	$1 + 0 \times 4$	$1$	$0^2 + 1^2$	是
第 2 项 ( $k=1$ )	$1 + 1 \times 4$	$5$	$1^2 + 2^2$	是
第 3 项 ( $k=2$ )	$1 + 2 \times 4$	$9$	$0^2 + 3^2$	是
第 4 项 ( $k=3$ )	$1 + 3 \times 4$	$13$	$2^2 + 3^2$	是
第 5 项 ( $k=4$ )	$1 + 4 \times 4$	$17$	$1^2 + 4^2$	是

全部 5 项均能在 $p, q \leqslant 7$ 的范围内表示为两平方和，因此 $(a=1, b=4)$ 是一个合法等差数列。