Cornacchia 算法 & P2063 二平方和定理

fangzichang · 2025-09-09 16:27:06 · 题解

闲话：在学校不但睡不好而且甚至睡不着，睡眠时间还被新来的校长砍了，而且我没得打 WA2，我甚至还没打完 coda 篇任意一条线路，周日没有休息，恐怕要初赛后才能打。

闲话：很难说这是一个特别的算法，还是一个欧几里得辗转相除的 trick，不过我似乎没有看到一个合适的中文介绍。不论如何，它很方便地解决了我在模拟赛中遇到的问题。

一道模拟赛题

假设我们要解决这样一个问题。

给定一个正整数 2\le n\le 10^{18}，构造一组解 (a,b) 满足 a,b 均为非负整数且 a^2+b^2=n，或报告无解。

本题有一些和 Cornacchia 算法完全无关的做法，复杂度从四次根号到几个 \log 都有，但是相比之下 Cornacchia 算法可能更为通用（以及优美）。
对于这个问题，如果使用该算法，主要瓶颈在质因数分解以及二次剩余上。

首先是一些简单的推导。

::::info[“高斯整数”？]{open} 以防有人像我一样不知道什么是高斯整数。

高斯整数是形如 a + bi 的复数，更准确地说，这篇文章中涉及的复数都是高斯整数。 ::::

大家知道，复数相乘时模长相乘。将问题转化一下，考虑去构造一个复数序列 w_{1\cdots k} 使得 |\prod_{i=1}^k w_i|^2=n，当然每个复数的实部和虚部都是非负整数。然后就只需要输出 \prod_{i=1}^k w_i 的实部和虚部就可以了，记得取绝对值。

当然，上面的讨论看起来一点信息量都没有。所以先质因数分解，设 n=\prod_{i=1}^{cnt} p_i^{c_i}，对于每个 p_i 分别构造。

1. 令 $v=1+\mathrm{i}$ 易得 $|v|^2=2$，那么往 $w$ 序列里插入 $c_i$ 个 $1+\mathrm{i}$，即可让 $|\prod w|^2$ 的值乘上 $2^{c_i}$。
2. - 对于 $c_i\equiv 0 \pmod 2$，往 $w$ 序列里插入 $c_i/2$ 个 $p_i$（注意是高斯整数 $p_i$，更精确地说是 $p_i+0\mathrm{i}$）即可，每个 $p_i$ 对 $|\prod w|^2$ 的贡献显然是 $p_i^2$； - 否则报告无解。 ::::info[为什么无解？] 学过小学奥数的朋友都知道，完全平方数模 $4$ 余 $0$ 或 $1$。 没学过也不用担心，学过抽象代数的朋友都知道，整数中模 $4$ 余 $3$ 的素数 $p$ 在高斯域中是惯性的，$(p)$ 仍然是素理想。 如果都没学过，打表找规律也不失为一个好办法。 你问我初等方法？因为 $p$ 是奇数，若 $p$ 是完全平方数，则 $n=\sqrt p$ 是奇数。设 $n=2k+1(k\in \mathbb N)$，有 $n^2=4k^2+4k+1\equiv 1 \pmod 4$。 ::::
3. 我断言，我们可以使用 Cornacchia 算法在 $\mathcal O(\log^2 V)$ 时间内构造一组 $(a_i,b_i)$ 满足 $a_i^2+b_i^2=p_i$，然后往 $w$ 里面插入 $c_i$ 个 $a_i+\mathrm{i}b_i$ 即可。根据费马平方和定理，存在满足条件的解 $(a_i,b_i)$。

当然了，实现的时候肯定要用快速幂把相等的复数先积起来，然后把所有东西乘起来。

它的正确性由过程给出。复杂度瓶颈，如前所述，在于写出高效的质因数分解和二次剩余。

P2063 二平方和定理

给定一个正整数 2\le n\le 10^{18}，求出所有非负整数数对 (a,b) 满足 a^2+b^2=n。

相比于上面那个题，这次我们需要构造所有的解。

可以发现，刚刚关于高斯整数的模型似乎仍然是适用的。
由唯一分解定理可以得到，最终仍旧得先质因数分解到每个质因子，再全部乘起来。
我们试着进一步分析看看，对于每个质数有多少种对应到高斯整数的方法，来试图遍历到所有解。

1. 仍旧仅有 $1+\mathrm{i}$ 一种可能，在实部和虚部分别枚举 $0,1,2$ 即可得到。
2. 与刚刚的分析相同，必须得把 $p_i$ 两两配对来产生若干个 $p_i^2$ 的贡献。用同样的手法可以证明，只有 $p_i+0\mathrm i$ 和 $0+p_i\mathrm i$ 两种可能的情况。
3. 类似的只有两组本质相同的解 $(a_i,b_i)$ 和 $(b_i,a_i)$（容易注意到 $a_i \ne b_i$）。 ::::info[引理，等一会还要用到]{open} 若 $(x,y)$ 是方程 $x^2 + d\cdot y^2 = p$ 的一组非负整数解，且 $y \not \equiv 0 \pmod p$，则 $\pm\sqrt{-d} \equiv x/y \pmod p$。 为了方便和我一样不会数论的朋友，写一下推导。 $$ \begin{aligned} x^2 + d\cdot y^2&\equiv 0 &\pmod p\\ x^2&\equiv -d\cdot y^2 &\pmod p\\ \frac{x^2}{y^2}&\equiv -d &\pmod p\\ \frac{x}{y}&\equiv \pm\sqrt{-d} &\pmod p\\ \end{aligned} $$ :::: ::::info[证明] 假设存在两组不同的非负整数解 $$(a, b)$$ 和 $$(c, d)$$，满足 $$a^2 + b^2 = p$$ 和 $$c^2 + d^2 = p$$，且 $$\{a, b\} \neq \{c, d\}$$。由于 $$p$$ 是质数，显然 $$a, b, c, d$$ 均为正整数。 考虑模 $$p$$ 下的运算。由 $$a^2 + b^2 = p$$， $$c^2 + d^2 = p$$，结合刚刚的讨论，$$a/b$$ 和 $$c/d$$ 都是模 $$p$$ 下 $$-1$$ 的平方根。由于 $$p \equiv 1 \pmod{4}$$，$$-1$$ 的二次剩余恰好有两个互为相反数的解，$$a/b \equiv \pm c/d \pmod{p}$$。 设 $$a/b \equiv c/d \pmod{p}$$，则 $$a d - b c \equiv 0 \pmod{p}$$。 由于 $$a, b, c, d < \sqrt{p}$$，显然 $$a d < p$$ 且 $$b c < p$$，所以 $$|a d - b c| < p$$。 因此 $$a d - b c = 0, a/b = c/d$$。 设 $$a/b \equiv -c/d \pmod{p}$$，容易得到 $$a d + b c \equiv 0 \pmod{p}$$。 由于 $$0 < a, b, c, d < \sqrt{p}$$，我们有 $$0 < a d + b c < 2p$$，所以 $$a d + b c = p$$。 发挥一下注意力，我们有： $$ \begin{aligned} (a d + b c)^2 + (a c - b d)^2&= a^2 d^2 + 2 a b c d + b^2 c^2 + a^2 c^2 - 2 a b c d + b^2 d^2\\ &= a^2 (c^2 + d^2) + b^2 (c^2 + d^2)\\ &= (a^2 + b^2)(c^2 + d^2)\\ &= p^2 \end{aligned} $$ 由于 $$a d + b c = p$$，我们有 $$p^2 + (a c - b d)^2 = p^2$$，因此 $$a c - b d = 0, a c = b d$$，我们得到 $$a/b = d/c$$，这与早些时候得到的 $$a/b = c/d$$ 似乎本质相同（$c,d$ 顺序无关），放在一起讨论。 设 $$k = a/b = d/c$$，则 $$a = k b$$ 且 $$d = k c$$。代入 $$a^2 + b^2 = p$$ 得到 $$k^2 b^2 + b^2 = b^2 (k^2 + 1) = p$$。类似地有 $$c^2 + k^2 c^2 = c^2 (k^2 + 1) = p$$。 因此，$$b^2 (k^2 + 1) = c^2 (k^2 + 1)$$，从而 $$b = c$$。于是 $$a = k b = k c$$ 且 $$d = k c = k b$$，所以 $$a = d$$，因此 $$\{a, b\} = \{c, d\}$$。 综上，$$p$$ 表示为两个平方数之和的表示在无序意义下是唯一的。 ::::

那么就简单了，每个质数写成高斯整数之后，只有一种本质不同的表示方法，顶多交换一下两个维度。直接对着搜索，枚举两类表示方法分别出现了几次，复杂度就是对的。

当然了，答案的两维也可以交换，如果担心漏解了，可以加上这一步，然后再去重。

除此之外再次强调，一堆复数乘起来可能有负的，记得取绝对值。使用 128 位整数来防止乘法导致的溢出。

给出一份参考代码，去除了宏定义与缺省源，它的语义应该比较明确，可以当作伪代码读。

时间复杂度，由于我数论不太行，可能只能分析到 \log^2 级别，不过仍旧是够用的，主要是我不知道一些数值应该用什么符号。据我观察，另一篇题解似乎与这个解法本质相同，复杂度也应该类似吧。

实践上跑得挺快。

struct cpx{
    i128 x,y;
    cpx(i128 p=0,i128 q=0){x=p,y=q;}
    void operator=(const i128&p){*this={p,0};}
    cpx operator*(const cpx&c)const{
        return cpx(x*c.x-y*c.y,x*c.y+y*c.x);
    }
};//复数类
template<class T>
T q_pow(T a,ll b){
    T res;
    for(res=1;b;b>>=1,a=a*a)
        if(b&1)res=res*a;
    return res;
}
i128 _abs(i128 a){
    return a>=0?a:-a;
}
void dfs(int i,cpx res,const vector<pair<cpx,int>>&vt,vector<cpx>&ret){
    if(i==(int)vt.size()){
        res.x=_abs(res.x),
        res.y=_abs(res.y),
        ret.push_back(res),
        swap(res.x,res.y),
        ret.push_back(res);
        return;
    }
    vector<cpx>p0;
    cpx it=vt[i],vl=1;
    for(int c=0;c<=vt[i].y;c++,vl=vl*it)
        p0.push_back(vl);
    swap(it.x,it.y),vl=1;
    for(int c=0;c<=vt[i].y;c++,vl=vl*it)
        //枚举：有 c 个高斯整数交换了两维，而 vt[i].y-c 个没有交换
        dfs(i+1,res*p0[vt[i].y-c]*vl,vt,ret);
}
pair<ll, ll> Cornacchia(const ll& P, ll p, ll t, ll d);//等下给出完整实现
signed main(){
    for(cin(T);T;T--){
        cin(n);
        vector<ll>d=PollardRho::factor(n);//质因数分解
        __gnu_pbds::gp_hash_table<ll,int>mp;
        for(ll&pr:d)mp[pr]++;
        if([&]()->bool{
            for(auto[v,t]:mp)
                if(v%4==3&&t%2==1)return 1;
            return 0;
        }()){
            puts("0");
            continue;
        }
        cpx res=1;
        vector<cpx>ret;
        vector<pair<cpx,int>>vt;
        for(auto[v,t]:mp){
            if(v==2)
                res=res*q_pow(cpx(1,1),t);//其实理论上直接加入 vt 也是对的，交换两个 1 不会得到错的解，但还是稍微精细一点好了
            else if(v%4==3)
                vt.push_back({cpx(v,0),t/2});
            else{
                ll I=Cipolla::getI(v);//求 v-1 也就是 -1 的二次剩余
                pair<ll,ll> w=Cornacchia(v,v,I,1);
                vt.push_back({cpx(w.x,w.y),t});
            }
        }
        dfs(0,res,vt,ret);
        auto cmp1=[&](cpx a,cpx b){
            return tie(a.x,a.y)<tie(b.x,b.y);
        };
        auto cmp2=[&](cpx a,cpx b){
            return tie(a.x,a.y)==tie(b.x,b.y);
        };
        sort(all(ret),cmp1),
        ret.erase(unique(all(ret),cmp2),ret.end());
        print(ret.size()),pc('\n');
        for(auto[x,y]:ret)
            print(x),pc(' '),print(y),pc('\n');
    }
    return 0;
}

Cornacchia 算法

进入正题。

只讨论经典情形。

设我们需要求方程 x^2+d\cdot y^2 = p 的一组解 (x,y)，其中 p 是奇素数且 p\nmid d，常见的特例是 d=1。
算法假定模方程 t^2\equiv -d\pmod p 存在解，并断言这是原方程有解的必要条件。当 d> 1 时这个条件并不是充分的，需要额外验证。

算法流程很简单。

1. 取一个合适的 t\in(0,p) 作为所有 t^2\equiv -d\pmod p 的解的代表。
   • 使用合适的二次剩余算法，比如 Cipolla 算法，可以很快地完成这一步。
   • 如果取了 t\in(0,p/2] 的较小的代表考虑，证明可能会稍微方便一点，但取哪个都不影响正确性。
2. 对 (p,t) 做辗转相除法，尝试得到一个余数序列 r，流程大概是 r_0=p,r_1=t,r_{i+1}=r_{i-1}\bmod r_i(i\ge1) 这样不断推下去，直到取到第一个足够小的余数 r_m，满足 r_{m}^2\le p。
3. 设 x=r_m，检验 (p-x^2)/d 是否为整数，然后检验是否为完全平方数。若是，则输出解 (x,\sqrt {(p-x^2)/d})，否则报告无解。

下面是一个可行的简短示例。

using ll = long long;
pair<ll, ll> Cornacchia(const ll& P, ll p, ll t, ll d) {
    if ((__int128) t * t <= P) {
        if ((P - t * t) % d != 0)
            return make_pair(-1, -1);
        ll y = (ll)(sqrt((P - t * t) / d) + .5);
        //为避免潜在精度误差，最好手动实现一个整数平方根算法
        if (y * y * d == P - t * t) 
            return make_pair(t, y);
        else return make_pair(-1, -1); 
    }
    return Cornacchia(P, t, p % t, d);
}

值得注意的是，虽然刚刚没怎么讨论符号，但是容易看出这个实现解出来的 x 和 y 都是非负的，我们有 0\le x,y\le \sqrt p。

即使只是对于 Cornacchia 算法，复杂度瓶颈也不在上述的辗转相除上，而是在求二次剩余上。

证明

受限于时间和水平，我只能给出大概的思路和前几步，如果你有兴趣，不妨试着完整证明一遍。

欧几里得算法有一个结论。假设我们在对 (p,t) 做辗转相除，将余数序列 r 拿出来，归纳可以证明，每个余数都可以写成 r_i=A_i\cdot p+B_i\cdot t 的形式，其中 A_i,B_i\in \mathbb Z。

::::info[r_i=A_i\cdot p+B_i\cdot t？] 为了方便和我一样不会数论的朋友，写一下解释。

对 i=0,1 可直接给出：r_0=1\cdot p+0\cdot t,r_1= 0\cdot p+1\cdot t。
对更大的 i，我们有 r_{i+1}=r_{i-1}\bmod r_i=r_{i-1}-\lfloor r_{i-1}/r_i\rfloor \cdot r_i。

当然也就可以递推给出