競プロ備忘録

競プロerの備忘録

CRT復元を書く

素数modで問題を解くためのアルゴリズムは色々有名なものがありますが、合成数modだとうまくいかないものも多いです。

このようなとき、素因数分解によって素数modの問題に分解することができれば、中国剰余定理を用いて元の問題の解を復元できることがあります。(いわゆるCRT復元)

AtCoderではC++を使っている限りはACLcrtという関数が用意されていますし、今は多くの言語にACLが移植されていますから、同等機能の関数なりメソッドなりで事足ります。
しかし、ある程度理解できるようになると案外ソラ書きできるなと気づいてきたので、メモしときます。

中国剰余定理について

概要です。証明はしません。

各文字についての条件は以下の通りです。

  •  N, i, p _ {i}は正整数
  •  a _ {i}は非負整数
  •  1 \le i \le N
  • すべての iについて a _ {i} \lt p _ {i}
  •  p _ {1}, p _ {2}, ..., p _ {N}は共通因数を持たない(すべての2つ組が互いに素)

これらの条件下で、 x \equiv a _ {1} \pmod {p _ {1}},  x \equiv a _ {2} \pmod {p _ {2}}, …,  x \equiv a _ {N} \pmod {p _ {N}}をすべて満たす正整数 xを見つけたいとします。
中国剰余定理は、このような x 0 \le x \lt p _ {1}p _ {2}...p _ {N}の範囲で一意に定まることを主張する定理です。

合成数modで問題を解きたいが、そのために便利なアルゴリズム素数modを要求している、という場合にこれが便利です。
素因数分解によって素数modの問題に分解し、それぞれが解をもつならば、中国剰余定理によって元の問題の解も復元できます。

あるいは、modを取れば効率的に求められるアルゴリズムがあって、これをmodを取らない場合の解に復元するのにも中国剰余定理が役に立ちます。
具体的には、いくつかの異なる大きな素数の法を設定して複数回問題を解きます。設定した法の総積が元の問題の解の上界を上回っているなら、中国剰余定理によって解を復元できます。

これらの解を復元する操作が俗にCRT復元とか言われるものです。

CRT復元のアルゴリズム

解きたい合同式は上記の通り何本あってもいいわけですが、2本の場合で解ければ繰返し適用するだけでよいので、合同式が2本ある場合のみを対象にします。

 x \equiv a _ {1} \pmod {p _ {1}}
 x \equiv a _ {2} \pmod {p _ {2}}

ここでまず、 a _ {1} = a _ {2}の場合、 x = a _ {1}が解になります。
したがって、以下は a _ {1} \ne a _ {2}としてよいです。

また、一般に a _ {1} \lt a _ {2}としてよいです。
そうでないとき、2本の合同式を入れ替えることで a _ {1} \lt a _ {2}にできます。

こんな感じの方程式はまずmodを使わない不定方程式に書き直します。
新しい整数 kを導入するとそれぞれ、

 x = a _ {1} + k _ {1}p _ {1}
 x = a _ {2} + k _ {2}p _ {2}

となります。2本の式で xは共通しているので、1本の方程式にまとめます。

 a _ {1} + k _ {1}p _ {1} = a _ {2} + k _ {2}p _ {2}

 kのついている項は左辺に、定数項は右辺にまとめると、

 k _ {1}p _ {1} - k _ {2}p _ {2} = a _ {2} - a _ {1}

ここまでわかると、2元1次不定方程式を解けばいいんだなとなります。条件から gcd(p _ {1}, p _ {2}) = 1かつ a _ {1} \ne a _ {2}なので、これは必ず解を持ちます。

2元1次不定方程式をとく

結論、拡張ユークリッドの互除法を書けばいいんですが、なんとなく拡張ユークリッドの互除法って ax + by = 1を解くイメージがあるので、2つ目の項が負になってると微妙に混乱しちゃいます。

先ほどの k _ {1}, k _ {2}についての不定方程式を眺めると、 p _ {2}でmodとってもよさそうだなと思えます。

 k _ {1}p _ {1} \equiv a _ {2} - a _ {1} \pmod {p _ {2}}

 a _ {2} \lt p _ {2}であって 0 \le a _ {1} \lt a _ {2}ですから、 a _ {2} - a _ {1} 0 \lt a _ {2} - a _ {1} \lt p _ {2}です。
ところで gcd(p _ {1}, p _ {2}) = 1なので、 p _ {1} \pmod {p _ {2}}は逆元を持ちます。
したがって、

 k _ {1} \equiv p _ {1} ^ {-1}(a _ {2} - a _ {1}) \pmod {p _ {2}}

うん。まあそれはそうかって結果なのですが、なんか見慣れた形になるので、安心して計算できます。

 p _ {1} ^ {-1} \pmod {p _ {2}}を求める

 p _ {2}素数ならみんな大好きフェルマーの小定理でなんとかできますが、今回はそういう制約ではないのでダメです。

結局のところ、拡張ユークリッドの互除法でなんとかします。
拡張ユークリッドの互除法をソラ書きするのは、noshiさんのブログを見るのが覚えやすいです。

以下の2本の自明な恒等式を考える、ということだけ覚えればよいです。

 p _ {1} = 1 \cdot p _ {1} + 0 \cdot p _ {2}
 p _ {2} = 0 \cdot p _ {1} + 1 \cdot p _ {2}

この式の両辺で、左辺でGCDを取る操作と同じ操作をすることを考えます。
 p _ {1} p _ {2}で割った余りをとるのは p _ {1} - \lfloor {\frac {p _ {1}}{p _ {2}}} \rfloor p _ {2}と同じなので、下の式の両辺に \lfloor {\frac {p _ {1}}{p _ {2}}} \rfloorをかけて上の式から引けばよいです。

これを繰り返すと、最終的に片方は左辺が 0になり、そうでないほうは以下のような形になります。

 gcd(p _ {1}, p _ {2}) = u \cdot p _ {1} + v \cdot p _ {2}

条件から gcd(p _ {1}, p _ {2}) = 1なので

 1 = u \cdot p _ {1} + v \cdot p _ {2}

と同じで、両辺 \mod {p _ {2}}で考えれば、

 1 \equiv u \cdot p _ {1} \pmod {p _ {2}}

となり、結果、 p _ {1} ^ {-1} \equiv u \pmod {p _ {2}}です。

実装例は以下のようになります。

Rust Playground

個人的には、再帰的なルーチンよりもこっちのほうが素直でわかりやすいと思います。
意味はないですが、一応上の式で言うところの v \mod {p _ {1}}での p _ {2}の逆元になっていることも確認しておきます。
ちなみに、もっとデカい数字でやると乗算でオーバーフローすることはないのかというのが心配ですが、しないらしいです。
なぜかはまだ理解してないので、とりあえず今はそれを信用することにしておきます。

CRT復元の実装の続き

 p _ {1} ^ {-1} \pmod {p _ {2}}はわかったので、これをつかって

 k _ {1} \equiv p _ {1} ^ {-1}(a _ {2} - a _ {1}) \pmod {p _ {2}}

を求めます。

なんかこれで満足しちゃいそうですが、求める対象は xだということを思い出します。
求めた k _ {1} \pmod {p _ {2}}を使うと、 xは以下のように表せます。

 x \equiv a _ {1} + k _ {1}p _ {1} \pmod {p _ {2}}

ところで、 a _ {1} \lt p _ {1}ですが、 k _ {1} \lt p _ {2}だとすると、

 a _ {1} \le p _ {1} - 1, k _ {1} \le p _ {2} - 1

なので、

 a _ {1} + k _ {1}p _ {1} \le p _ {1} - 1 + (p _ {2} - 1)p _ {1} = p _ {1}p _ {2} - 1 \lt p _ {1}p _ {2}

となります。よって、 k _ {1} \pmod {p _ {2}}を計算する時点で正しく 0 \le k _ {1} \lt p _ {2}に収まっていれば、求めた k _ {1}をそのまま使ってよいことがわかります。

したがって、

 x = a _ {1} + k _ {1}p _ {1}

を計算することで xが求まります。

実装例は以下の通りです。

Rust Playground

Wikipediaの「中国の剰余定理」のページに載っていた例ほぼそのままですが、上手くいってそうでしょうか。

 gcd(p _ {1}, p _ {2}) \ne 1の場合

ここまで、 gcd(p _ {1}, p _ {2}) = 1の場合のみを考えてきましたが、そうでない場合にも拡張ができます。
主張や証明についてはけんちょんさんのQiitaの記事とか高校数学の美しい物語とか見るのが良いと思いますが、雑に言えば p _ {i}の総積ではなくて、すべての p _ {i}のLCMまでの範囲で一意に xを決めることができるようになります。
ただし、すべての p _ {i}が互いに素な場合と異なり、解が存在するための条件が加わります。

また2つの合同式から、1本の方程式にまとめます。

 k _ {1}p _ {1} - k _ {2}p _ {2} = a _ {2} - a _ {1}

ここで、 gcd(p _ {1}, p _ {2}) \ne 1なので、それを gと置くと、左辺は gでくくることができます。
したがって、右辺が gの倍数でない場合には解がありません。つまるところ、

 a _ {1} \equiv a _ {2} \pmod {gcd(p _ {1}, p _ {2})}

が解が存在する条件となります。

 p _ {1} = g \cdot p' _ {1}, p _ {2} = g \cdot p' _ {2}とすると、先ほどの方程式は

 k _ {1}p' _ {1} - k _ {2}p' _ {2} = \frac{a _ {2} - a _ {1}}{g}

となり、 gcd(p' _ {1}, p' _ {2}) = 1ですから、これは先ほどと同じように解けます。

 k _ {1} \equiv {{p' _ {1}} ^ {-1}} \cdot \frac{a _ {2} - a _ {1}}{g} \pmod {p' _ {2}}

となります。

さて、ここから x = a _ {1} + k _ {1}p _ {1}を用いて xの具体的な値を求めたいですが、 k _ {1}はどのような値であるべきでしょうか?

とりあえずは k _ {1} = ({{p' _ {1}} ^ {-1} \mod {p' _ {2}}}) \cdot \frac{a _ {2} - a _ {1}}{g} \mod {p' _ {2}}としておきます。
ところで、少なくとも、以下の等式は正しいことがわかります。( rは適当な整数)

 p' _ {1} \cdot ({{p' _ {1}} ^ {-1}} \mod {p' _ {2}}) + r \cdot {p' _ {2}} = 1

両辺に gをかけると、

 p _ {1} \cdot ({{p' _ {1}} ^ {-1}} \mod {p' _ {2}}) + r \cdot {p _ {2}} = g

これを \mod {p _ {2}}で評価すると、

 p _ {1} \cdot ({{p' _ {1}} ^ {-1}} \mod {p' _ {2}}) \equiv g \pmod {p _ {2}}

です。
これを用いて、 x = a _ {1} + k _ {1}p _ {1}の右辺に k _ {1}の値を代入した結果が x \equiv a _ {1} \pmod {p _ {1}}, x \equiv a _ {2} \pmod {p _ {2}}を満たすか確認します。

まず、 \mod {p _ {1}}で評価すると、 k _ {1}p _ {1}の項が消えるので、条件を満たします。

次に、 \mod {p _ {2}}で評価すると、

 a _ {1} + k _ {1}p _ {1} \equiv a _ {1} + ({{p' _ {1}} ^ {-1} \mod {p' _ {2}}}) \cdot \frac{a _ {2} - a _ {1}}{g} \cdot {p _ {1}} \pmod {p _ {2}}

ここで、 p _ {1} \cdot ({{p' _ {1}} ^ {-1}} \mod {p' _ {2}}) \equiv g \pmod {p _ {2}}なので、さらに変形すると、

 a _ {1} + {p _ {1}} \cdot ({{p' _ {1}} ^ {-1} \mod {p' _ {2}}}) \cdot \frac{a _ {2} - a _ {1}}{g} \equiv a _ {1} + g \cdot \frac{a _ {2} - a _ {1}}{g} \equiv a _ {1} + a _ {2} - a _ {1} \equiv a _ {2} \pmod {p _ {2}}

となり、条件を満たします。

よって、 k _ {1} = ({{p' _ {1}} ^ {-1} \mod {p' _ {2}}}) \cdot \frac{a _ {2} - a _ {1}}{g} \mod {p' _ {2}}として問題なく、 x = a _ {1} + k _ {1}p _ {1}が解となります。

 0 \le x \lt lcm(p _ {1}, p _ {2})に収まっているかどうかについては、 0 \le k _ {1} \lt p' _ {2}であることからわかります。

実装例は以下の通りです。

Rust Playground

ext_gcdがGCDを追加で返すように関数定義を変更し、GCDが1でなくてもパニックしないようにassert!を削除しました。
こうして返されたGCDが1より大きいか否かで処理を分岐しています。(実際はよく見るとわかる通り、ルーチンを分岐する必要はない。)
crt(1, 8, 5, 12)で、8と12は互いに素ではないですが、ちゃんとLCMである24までの間に満たす値を発見して返しています。

まとめ

最初はソラ書きすることが目的だったんですが、だんだん話がソラ書きとは関係ない方向にそれたので、単に書き方の学習メモになりました。

一応、上記全部ソラでなんとか書けるようになりましたが、以下のことを最低減覚えておけばその場でアルゴリズムが導出できるとわかりました。

  • ext_gcdの書き方:2本の恒等式を書くことと、左辺でGCDを取る操作をすることだけ覚えればよいです
  • 合同式と等式の変換:法を何か整数とかけて足せば合同式は等式になるし、2元1次不定方程式はどっちかの係数でmodを取れば合同式にできます

あとは出てきた式をごにょごにょいじり倒せば、その場で何とかできます。
とはいえ、なんとなく雰囲気で流していた部分もあるので、記事書きながら理解を深められて良かったです。