2024-03-29

CRT復元を書く

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

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

AtCoderではC++を使っている限りはACLにcrtという関数が用意されていますし、今は多くの言語に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元１次不定方程式はどっちかの係数でmodを取れば合同式にできます

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

2024-03-17

ABC345E - Colorful Subsequence

3完しかできなくてひどい目に遭った…
ふと思い返すと、ひどい目に遭ったときしかブログ書いてない気がする。

そんなことはどうでもよくて、Dは解けそうにないし面白い要素もないので、Eをずっと考えていた。
こっちは考えごたえがあってとても面白かったので、コンテスト後だけど自力ACできて嬉しい。

解法

とりあえず愚直なDPを考えてみる。 $K \lt 500$ や実行時間制限5secという不自然な制約があり、ワンチャン通るかもしれないので、書いてみる価値は十分ある。
DPの定義は以下のようになる。

$DP _ {i, j} :=$ 　前からボールの使い方を決めることとして、 $i$ 番目を使い、そこまでで $j$ 個ボールを飛ばしている場合の最高スコア

愚直にDP配列を埋めるのは簡単で、前からボールを見つつ、色が異なる自身の直前 $K$ 個以内のボール $prev$ にたいして、 $DP _ {i, j} = max( DP _ {i, j}, DP _ {prev, j - (i - prev - 1)} + P _ {i})$ と更新していけばよい。
この定義だと、 $DP _ {i, j}$ では $i$ が必ず使われることになっているので、最初と最後の場合分けが面倒くさい。なので、最初と最後に番兵として $(C, P) = (0, 0)$ などを挟むとよい。

時間計算量は $O(NK ^ {2})$ となる。一応投げてみると、半分以上のケースでTLEすることがわかる。（提出例：Submission #51323752 - Monoxer Programming Contest 2024（AtCoder Beginner Contest 345））

コードの見た目的に、for prev in (0..i).rev() {....の中身を小さい計算量でできそうなので、ここの高速化を考える。

dp[i][j].max(dp[prev][j-skip] + p[i].1);という箇所に着目すると、p[i].1はずっと同じなので考える必要はない。dp[prev][j-skip]という箇所を考えると、let skip = i - prev - 1;なので、dp[prev][j - (i - prev - 1)]と同じ。この2つの添え字の差を考えると、 $prev - j + (i - prev - 1) = i - j - 1$ であり、prevに依存しなくなることがわかる。
これにすぐ気づくのは難しかったが、DP配列のある箇所の更新の一連の操作を紙でマス目に書いてみると、マス目をナナメに走ることから何となくインデックスの和か差のどっちかが重要なんだろうということが予想できる。

ここまでわかると、DP配列のナナメ方向の累積maxを持てれば良さそうだと予想できる。
ただ、色の制約があるので、実際には累積スコアと最後に使うことを決めたボールの色のペアを、スコアの上位2つ分色違いになるように持てば良い。（以後も、このデータの持ち方やこれを維持する操作を指して「累積max」「max」ということにする）

累積maxということなので、セグメント木が真っ先に思い浮かぶ。しかし、これを書くとTLEする。（提出例：Submission #51353665 - Monoxer Programming Contest 2024（AtCoder Beginner Contest 345））
VecDeqでセグメント木を持つ数を減らしているが、これをやらないとMLEもする。
計算量は $O(NKlogK)$ なので通るだろうと踏んでいたが、よく考えると $NKlogK$ は最大 $10 ^ 9$ 近くなるので、まあ通らなくても文句は言えない。

さらにこのコードを睨むと、累積maxを取得する箇所は、常に $0$ から $j$ までの累積maxを取得していることがわかる。
中途半端な区間の累積maxが不要であるなら、わざわざセグメント木を使うまでもなく、単に配列で保持することができる。累積maxの更新順序に着目すると、jは小さいほうから更新されるため、値更新を行ったら直前の要素とmaxを取るだけで累積maxの更新もできることがわかる。

ここまでやると、無事ACすることができる。計算量は $O(NK)$ 。
提出例は以下の通り。

Submission #51356969 - Monoxer Programming Contest 2024（AtCoder Beginner Contest 345）

2024-02-19

AHC30A - Polyomino Mining

基本マラソンはあまり真剣にやらないんですが、今回はできそうなことがありそうなのでちょっと時間をかけてみました。

暫定200位よりちょっと悪いくらいらしいので、私の実力にしてはよくやったほうだと思います。
私は焼きなましだとかビームサーチだとか、マラソン界隈でよく耳にするアルゴリズムは一切書けないし、統計に関する知識もほぼゼロなので、そういう知識を使わない解法になります。

提出順に書いていますが、改善がないやつとかパラメータ調整程度の提出は省略します。
スコアは絶対スコアです。また、N投目のNは間違えてる可能性あります。WAやTLEは除いています。

1投目、12696000000点

サンプルプログラムを投げるだけ。私はPythonアンチなので、Rustに書き直して提出した。（というのは半分嘘で、サンプル見る前に書いてたプログラムがたまたまサンプルと同じアルゴリズムだっただけ）

3投目、10554000000点

明らかにどの油田も来ないような位置は全探索で容易にわかるので、それを調べないようにする。

4投目、7768000000点

端っこから順に掘るのはなんか無駄そうなので、順番を工夫する。
この後何回かパラメータは調整したが、この時点では、配置されうる油田の数が多いマスを優先し、あとはなるべく隣り合うマスを連続で調べないようにした。

5~9投目、6950000000→6299000000点

探索順の調整を頑張る。

油田があるマスもないマスも割と集約して位置するような気がするので、周囲の油田がないことが確定しているマスの数に合わせて適当なペナルティを振ることにする。油田があるマスの周りに油田があることは多いと思うので、周囲に油田があるならその分はペナルティを減らしてあげる。

ペナルティの振り方はよくわからなかったが、周囲の油田がないマスの数の5倍から油田があるマスの数を引いた数の平均が良さそうだったので、これでソートしてペナルティの少ないほうから探索した。

また、配置が確定した油田は最初に掘るように工夫したのもこの時点。

10投目、5755000000点

正直コードを見直しただけでは9投目との違いが何か判らなかったが、バグフィックスをしたら点数が上がったらしい。

12投目、5544000000点

よく考えると、M=2なら配置をすべてメモできるので、ダメな配置をより細かく把握できる。
なので、M=2のときだけやりかたを変える。

やりかたは単純で、全配置をメモして、1点掘るごとに現時点で可能な配置を全探索してダメな配置をリストから排除する、をひたすら繰り返す。

16投目、5373000000点

M=3のときも、頑張って高速化すれば全状態列挙ができるので、7投目のときと同じように全探索ベースで頑張る。

17投目、5352000000点

これも11投目と何が違うのかよくわからないが、同様にバグフィックスや高速化で点数が上がったらしい。

18投目、4783000000点

M>3のときでも、状態数が十分小さいことがわかっていれば全探索ベースでいけるので、あり得る状態数が一定の閾値を下回った時点で全探索ベースに切り替えるようにする。

たまに、掘っても掘っても状態数が小さくならないような場合があり、このときは下手するとTLEするので、タイマーで一定の時間を上回った時点で元のルーチンに戻ってくるようにしている。

19投目以降、最終4344000000点

基本的にはパラメータ調整しかしていない。

最後の方は、M>3の場合にいかに状態数を効率よく減らして全探索ベースのルーチンに持ち込めるかを考えていたが、あまり良い手は浮かばなかった。

全探索ベースのルーチンを高速化して閾値を上げられるようにすればスコアは上がるだろうが、具体的な高速化の方法はわからなかった。

その他

結局、eの使い方は最後までわからなかったので、読み捨てて何にも使いませんでした。
複数掘削で油田がない位置を効率的に探し当てられないかとか考えましたが、eが相当小さくないと探し当てるのは無理だなという結論になりました。

全探索ベースの解法の高速化の方法がわからなかったのも痛かったです。フローとか使えるんかなーとか、ビット演算で頑張れないかなーとか、雑なことは考えていましたが、効果的なものはわかりませんでした。

2024-02-11

ABC340E - Mancala 2

冒頭からネタバレですが、やたらとdiffが低かった割には想定解法が遅延セグメント木だったので、おったまげた人が多かったようです。私もコンテスト中は遅延セグメント木使いましたが、E問題ごときでそんなん想定解のはずないだろって思い込んでいたので、たまげました。
自分が緑くらいのときは遅延セグメント木なんて知らなかったと思うので、遅延セグメント木以外で解く方法も必要でしょう。

そんなわけなので、あえて想定解と違う方法でも解いてみたので記事にします。

解法

平方分割ができそうな感触があるので、平方分割を検討してみる。

この問題が解けるには、区間加算と一点の値取得ができればよい。
遅延セグメント木を使えばこれを両方 $O(\log N)$ 程度の計算量でできるが、平方分割の場合は $O(\sqrt N)$ で捌くことを目指す。

具体的な方法としては、箱を $\sqrt N$ サイズのバケットに分割し、 $\sqrt N$ サイズの累積和をとれる配列を $A$ とは別にバケットごとに持つことにする。
区間加算は、imos法によって各バケットの加算したい区間の端に足し引きすることで行う。一点取得は、欲しい値のあるバケットの累積和を全部計算して、 $A$ に加算することで行う。

まず区間加算の計算量を考察する。バケットごとの区間加算は、区間の端を足し引きするだけなので、 $O(1)$ で行える。
区間加算で足さなければいけないバケットは $O(\sqrt N)$ 個あり、ハンパを処理するのに追加で $O(\sqrt N)$ 個加算が必要になりうる。
したがって、クエリごと $O(\sqrt N)$ 程度の計算量で処理することができる。

次に一点取得の計算量を考察するが、これはimos法によって押し付けられた累積和の計算と累積和配列のクリアだけなので、 $O(\sqrt N)$ で行えることが容易にわかる。

したがって、全体で $O(M\sqrt{N})$ の時間計算量ですべてのクエリを捌けることがわかる。

実装例は以下。
Submission #50189546 - KAJIMA CORPORATION CONTEST 2024（AtCoder Beginner Contest 340）

36-42行目のところもループを回して $A$ に直接加算するようにしても、クエリごとの計算量は変わらず $O(\sqrt N)$ のはずだと思っているが、どうもここをサボるとTLEするらしい。
$N$ も $M$ もそれなりにデカいので $O(M\sqrt{N})$ が落ちても文句は言えないが、Rustくらい速い言語でも定数倍で落ちることがあるんだなあと思ったり思わなかったり…
なんかバグらせているだけかもしれないが、これ以上は追求しない。

あとがき

ちなみにコードが汚いのは平方分割書き慣れてないからです（というかまともに平方分割を使って問題解いたのは多分初めて）。
どうやったらもうちょいキレイに書けるんですかね？

（追記）一応補足しておくと、wrapping_add/wrapping_subの部分は普通にadd_assign/sub_assignで書いても良いし、その方が若干スッキリはする。が、ローカル環境でも常にrelease実行しなきゃいけなくて面倒なので、どっちがいいか問題はある。debug実行するとあらゆるところでパニックすることになる。
Submission #50192312 - KAJIMA CORPORATION CONTEST 2024（AtCoder Beginner Contest 340）

2024-02-09

技術室奥プログラミングコンテスト#6 Day1E - And Xor Or

再帰がハマってキレイに書ける問題は楽しいです。

Rustはスライスで自然に範囲を狭めていくような再帰が書けるのが嬉しいですね。

解法

AND,XOR,ORはそれぞれ&, ^, |で書くことにする。

数式のAND,XOR,ORの部分がなんか怪しいので、とりあえず $0,1$ を当てはめて真理値表を作ってみる。
すると、実はこれが $A _ {i}$ ^ $A _ {j}$ に等しいということがわかる。

足し算の部分は上手いこと変換できそうにないので、この式をもとに最大化の方針を考える。
定石として、こういうのはビットごとに考えると良く、また上位ビットほど結果に与える影響は強いから、上のビットをいい感じにしたい。

$A _ {i}, A _ {j}$ ともに1ビットしかないとして結果を観察すると、

$0 + 0 - (0$ ^ $0) = 0$
$0 + 1 - (0$ ^ $1) = 0$
$1 + 0 - (1$ ^ $0) = 0$
$1 + 1 - (1$ ^ $1) = 2$

となっている。まとめれば、 $(1, 1)$ の組み合わせのときは得だが、それ以外のときは得も損もしない。
さらに言えば、 $(1, 1)$ 以外では、ビットが $0$ なのか $1$ なのかを無視してよい。

したがって、以下のような方針で解ける。

$A$ の最大値のMSBを調べる。これと同じビットが立っている要素の集合を $S$ とする。
$S$ の要素数が2以上の場合は、 $S$ の間で組み合わせるのが最適なので、MSBをクリアして $A = S$ として同じ問題を解く。
そうでない場合は、MSBは結果に影響を及ぼさないので、クリアしたうえで再帰的に $A$ について問題を解く。

計算量について考察すると、各ビットの検査は最大値を求めて要素を抜き出すだけなので $O(N)$ で完了できることがわかり、再帰は $A$ の最大要素のビット数に左右されるので、結果的には $O(N \log max(A))$ となる。

実装例は以下。
Submission #50103292 - 技術室奥プログラミングコンテスト#6 Day1

ちなみに、無駄にもう1つlogを付けてよければ、以下のようにもう少しスッキリした実装にもできる。そして、なぜかこっちのほうが速いorz
Submission #50101708 - 技術室奥プログラミングコンテスト#6 Day1

Fastestにあと1msecと迫る速さだし、単純なルーチンなので非再帰化してみようと思って手を出したのが以下。
Submission #50101811 - 技術室奥プログラミングコンテスト#6 Day1

あと1msecがなかなか縮まらない…

2024-02-05

ABC339G - Smaller Sum

こんなん乗せられていいはずないだろ！
っていうのがセグ木(みたいな何か)に乗るの、好きです。

解説見て、これちゃんとした名前のついてるデータ構造だったのか～と思いました。

解法

問題を見たところこれに似てるなと感じた。
B - 買い物 2 (Shopping 2)

そうなるとセグ木に乗せたいなという気になる。
どう乗せるか考えてみると、区間ごとに $X$ 以下の要素の累積和が小さい計算量で求められると嬉しいので、区間ごとに $A _ {i}$ でソートされた配列の累積和を持つことにすると上手くいきそう。

これを愚直に実装するとどうなるか考える。
構築は2つの子ノードの配列を単純に結合してソートしたものを全ノード作り切ってから、改めて全ノード累積和を取ればよい。
クエリへの回答は、セグ木と同じようにボトムアップでノードをなめつつ、各ノードで二分探索によって $X$ 以下の累積和を足し込めばよい。

クエリへの回答の計算量を雑に評価すると、もともと $O(logN)$ で済んでいたセグ木の操作について、各ノード $O(logN)$ で値を取得する必要があるので、クエリごと $O(log ^ 2 {N})$ 、全体で $O(Q log ^ 2{N})$ となる。

構築のほうがなんかヤバそうな雰囲気になるが実は愚直のままでよい。
まず木の各段の要素数を考えると、( $A _ {i}$ をちゃんとマージすると若干減るが、)どの段も常に $N$ 個になっている。これはある段でノードが $L$ 個あり、各ノード $M$ 要素持っていたとして、親の段はそれぞれのノードが $2M$ 要素持つようになるが、ノードの個数が $\frac{L}{2}$ 個しかなく、段全体の要素数が変わらないことからわかる。
そして、各段の要素が1つ上の段に引き上げられるのはぞれぞれ1回ずつなので、木を構築するところまでは $O(NlogN)$ の時間計算量で済む。
そのあと全ノードでソートが必要になるが、これは最上段のソートが最も重くなるので、結局全体では $O(N log ^ 2{N})$ の計算量で構築が可能。

よって、木の構築とクエリへの回答を合わせて、全体で $O(Qlog ^ 2{N} + Nlog ^ 2{N})$ の計算量で済む。

実装例は以下。
Submission #49965872 - Japan Registry Services (JPRS) Programming Contest 2024 (AtCoder Beginner Contest 339)

あとがき

解説を見てみるとmerge-sort treeとか書いてあるんで改めて考えてみると、確かにこれって木のノードがマージソートする過程になっているわけですね。
つまり、木の構築の過程で各ノードをバカ真面目にソートする必要はなくて、ソート済配列のマージで済みます。木の構築の計算量が $O(NlogN)$ に減るのでうれしいです。

まあそんなこと知らなくてもソート済配列のマージの時点で気付けよって話なんですが、本番これ解き始めた時点で20分しか残ってなかったんで焦っていたんですよね…
致し方なしです。

2024-02-03

競プロ用Rust環境構築

記事で載せてる実装例とかほぼ全部Rustのコードなことからわかるように、私はRust使いです。(私の記事読んだことある人なんでほぼいないでしょうが)

たまに、Rustで競プロやりたいけど環境構築の方法がよくわからん、みたいな人もいるので、お手軽環境構築していきましょう。

リッチな環境じゃないですが、私はこの環境で出ててあんまり不便感じてないので、まあ大丈夫なはずです。

WSLをインストールする

のっけからこれかよって感じですが、WSLを入れていきましょう。
別にWindowsでも当然Rustは使えるんですが、私はWindowsでのプログラミングのことをよく知らないので、WSL入れていきます。

当たり前ですが、生のLinux使ってるならこの手順は不要です。

手順は簡単で、管理者としてPowershellを立ち上げてwsl.exe --installすればよいです。

必要に応じて以下を参照しましょう。ここに書いてある通り、よっぽど古い環境使ってるとかでなければwsl.exe --installで入るはずです。
WSL のインストール | Microsoft Learn

インストールされるのはUbuntuなので、Ubuntuが嫌いな人は別のディストロ入れればいいと思います。
wsl.exe --list --onlineすればわかりますが、Debian, Oracle Linux, Kali Linux, Open SUSE, SLESとか使えます。

build-essentialを入れる

sudo apt install build-essentialとすれば入ります。ほかのディストロにも似たようなのがあると思います。

なんでこんなの入れるの？って話ですが、後で導入するRust Analyzerがerror: linker 'cc' not foundなどというエラーで死ぬことがあるからです。

gcc入れるだけで十分な気もするのですが、なんかトラブっても嫌なのでとりあえず入れときます。

rustupする

LinuxでのRust環境の導入は超簡単で、rustupというツールを実行すれば勝手に入ります。

Rust をインストール - Rustプログラミング言語

たぶんcurl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | shとかいうコマンドが書いてあると思うので、これをそのまま黒窓にコピペして実行します。
今はこれが書いてありますが、将来もそうかはわからないので、ちゃんと公式ページからコマンドを拾ってきてください。

curlがないよって怒られたら、sudo apt updateして、sudo apt install curlでインストールしてから実行しましょう。

実行すると1回だけ入力を求められると思いますが、これはそのままEnterを叩けばdefaultでインストールを進められます。

インストール後、シェルを再起動するか、あるいは出力の末尾に出てくるsource "$HOME/.cargo/env"をそのままコピペして実行すれば、パスが通ります。

これで、cargo, rustc, rustfmt, rustupあたりのツールが入っているはずなので、確認してみましょう。command -vかwhichで、それぞれのツールのインストール場所のパスが見えればOKです。
私は今dockerでテキトーに立てたコンテナで試しているので、/root配下に置かれています。わざわざsudoでrustupしたりしてなければ、ホームディレクトリ配下にインストールされることになるでしょう。

$ command -v cargo
/root/.cargo/bin/cargo
$ command -v rustc
/root/.cargo/bin/rustc
$ command -v rustup
/root/.cargo/bin/rustup
$ command -v rustfmt
/root/.cargo/bin/rustfmt

手間でなければ一応バージョンも見ておきます。

$ cargo --version
cargo 1.75.0 (1d8b05cdd 2023-11-20)
$ rustc --version
rustc 1.75.0 (82e1608df 2023-12-21)
$ rustup --version
rustup 1.26.0 (5af9b9484 2023-04-05)
info: This is the version for the rustup toolchain manager, not the rustc compiler.
info: The currently active `rustc` version is `rustc 1.75.0 (82e1608df 2023-12-21)`
$ rustfmt --version
rustfmt 1.7.0-stable (82e1608 2023-12-21)

補完の設定をする

オプションです。おそらく現時点ではcargoを入力した後、タブ補完とか試しても何も反応しないでしょう。
これでは不便なので、シェルの補完をやっていきます。

もしbash-completionがインストールされていない場合は、sudo apt install bash-completionなどでインストールしておきましょう。
たしかWSLでは必要なかったと思いますが、念のためです。

rustup completionsコマンドを叩くとめっちゃ長いヘルプが出てきますが、その中にシェルごとの補完の設定の仕方が書いてあります。
WSLを何もカスタムしていなければbashを使っているでしょうから、BASH:と書いてあるところを読みながらコマンドを打ち込みましょう。

今時点では、以下のようなコマンドを打ち込めと書いてあります。

$ mkdir -p ~/.local/share/bash-completion/completions
$ rustup completions bash >> ~/.local/share/bash-completion/completions/rustup
$ rustup completions bash cargo >> ~/.local/share/bash-completion/completions/cargo

このあと1回ログアウトせよ、と書いてあるので、シェルを立ち上げ直します。
ちなみにdockerで試しているとき、ログアウトせずに補完を試したらbash: _get_comp_words_by_ref: command not foundなるエラーが出ていました。
これはsource /etc/bash_completionで治ります。 kubectl bash completion doesn't work in ubuntu docker container - Stack Overflow

シェルを立ち上げ直したら、補完を試しましょう。rustupコマンドも同様に補完が効くようになっているはずなので、手間でなければ試してください。

$ cargo <TAB>
~出力例は省略。上手くいっていればadd, check, clean, buildなど多数のコマンドが並んでいるはず~

エディタを入れる

なんでもいいです。VimでもEmacsでも好きなものを使うと良いでしょう。
私はVisual Studio Codeしか知らないので、その前提で話を進めます。

Visual Studio Code - Code Editing. Redefined

Visual Studio Codeをインストールすると、WSLからcodeコマンドが使えるようになるはずです。
Windowsから起動するのではなくWSLのシェルからcode <target directory>で起動することで、勝手にVisual Studio CodeがWSLに接続するようになります。

拡張機能を入れる

もちろん好きなものを自己責任で入れればよいですが、とりあえずRust Analyzerというやつは入れましょう。
一応リンクは貼りましたが、Visual Studio Codeの左端にある、右上の欠けた田のマークからインストールするのが楽です。

rust-analyzer - Visual Studio Marketplace

私は毎回環境立てるときは以下のRust Extension Packを入れてますが、特にどっかが公式で出してるわけではないので、入れる場合は自己責任で入れてください。
Rust Analyzerもここに入ってます。

Rust Extension Pack - Visual Studio Marketplace

コンテスト用のクレートを作る

Rustではクレートという単位でプロジェクトを分けます。依存モジュールもクレート単位で解決されます。
大きくバイナリクレートとライブラリクレートがありますが、混在できます。また、1つのクレートに複数のバイナリを持たせることもできます。

競プロやるにあたっては、例えば1つだけバイナリクレートを作って毎回上書きして書き捨てるとか、問題1つにつき1クレートとか(さすがにいないかな…)、色々やり方はあると思います。
私は1クレートで1コンテストにしているので、その方法でやります。

まずcargo newで新規クレートを作りましょう。例えば、

$ cargo new abc300

で、abc300というディレクトリが作られます。あるいは、mkdir abc300; cd abc300; cargo initでも同じことになります。
作成されたabc300は、以下のような構造になっているはずです。(もしかしたら.gitignoreがあるかもしれませんが、些細な差です)

$ tree abc300/
abc300/
|-- Cargo.toml
`-- src
    `-- main.rs

1 directory, 2 files

コンテスト用クレートが出来たら、これをcode abc300で開いてみましょう。

開いたウィンドウで、ターミナル > 新しいターミナル　としてターミナルを開いておきます。
そして、

$ mkdir src/bin -p
$ cd src/bin
$ touch {a..g}.rs
$ code *

として、コンテストで使うソースファイルを作ります。

src/main.rsはどうしたんだって話ですが、私は使いません。なので、別に消してもいいですし、残しておいても無害なのでそのままでも良いです。

外部クレートを入れる

RustでもAtCoderであれば外部ライブラリがたくさん使えるので、以下を参照しながら好きなものを入れていきましょう。
（スプレッドシートが真と書いてあるので、スプレッドシートを見たほうが良いかもしれない…）

https://img.atcoder.jp/file/language-update/language-list.html

外部ライブラリのダウンロードにはcargo addコマンドを使うと良いです。
例えばRust競プロer御用達のproconioをインストールするには、カレントディレクトリがabc300配下であるのを確認して

$ cargo add proconio@=0.4.5 --features "derive"

とすることで依存クレートに加えることができます。
--featuresは外部クレートのオプションの機能を指定するのに使います。proconioには出力高速化のためにfastoutというマクロが用意されていますが、これは--features "derive"を指定したときのみ使えます。

デフォルトではcrates.ioと呼ばれるレジストリからクレートはダウンロードされますが、ローカルフォルダやGithubのリポジトリからダウンロードすることも可能です。
詳しくはドキュメントを読んでみてください。
cargo add - The Cargo Book

ちなみに、cargo addの操作は、さっきクレートのディレクトリ構造を見たときにチラッと出てきたCargo.tomlの[dependencies]という項目にエントリを加えるのと同じです。なので、手書きでCargo.tomlを書き換えることでも外部クレートの追加は可能です。

コードの実行

あとはコードを書いて実行するだけです。

例えば、src/bin/a.rsを実行するときには、以下のようにすれば実行できます。

$ cargo run --bin a

また、cargoにはreleaseビルドとdebugビルドがあります。デフォルトではdebugビルドで、--releaseを指定することでreleaseビルドになります。

$ cargo run --bin a --release

だいたいdebugビルドでも困りませんが、AtCoderの環境はreleaseビルドであるということだけ覚えておいたほうが良いでしょう。

また、cargo newで作成されるsrc/main.rsを実行する場合には、クレート名を--binの引数にする必要があります。

$ cargo run --bin abc300

src/bin配下のバイナリがない場合には、--binがなくてもsrc/main.rsが実行されます。

$ cargo run

ここまででコンテストに出るには十分な環境が整いました！あとの章はオマケです。読まなくても支障はないです。

自作のライブラリを使う

C++とかだとライブラリを展開してバンドルするツールを使ってるのをよく見ます。Rustにもそういうツールはあります。

いくつか種類はあるのですが、私はcargo-equipというツールを使っています。

GitHub - qryxip/cargo-equip: A Cargo subcommand to bundle your code into one `.rs` file for competitive programming

crates.ioに公開されているので、cargo install cargo-equipで入手することができます。

使い方は簡単で、cargo equip --bin a > run.rsとして対象のバイナリが依存するライブラリを上手いことバンドルしたコードを標準出力に吐き出させ、リダイレクトで別のファイルに流せば、提出できる形になります。

ただ、何もオプションを付けないと森羅万象がバンドルされて大変なことになります。私が使うときは以下のオプションを付けて、AtCoderで使えるクレートを除外したり、cargo checkをパスしたり、minifyするオプションを付けてコード量を減らしたりといったことをしてもらっています。

cargo equip --exclude-atcoder-crates --remove docs --remove comments --exclude-atcoder-202301-crates --minify libs --no-rustfmt --no-check

長すぎるわ！って話ですが、cargo-equipは仕様上いっぺんビルドが走るので、ちょっとでも処理を減らさないとバンドルに許容できない時間がかかります。(Rustのビルドは遅いことで有名)

なので、.bashrcにequipというエイリアスで、上記のオプションを全部ひっくるめて登録しています。

$ alias equip
alias equip='cargo equip --exclude-atcoder-crates --remove docs --remove comments --exclude-atcoder-202301-crates --minify libs --no-rustfmt --no-check'

これでequip --bin a > run.rsのように、楽にオプションもりもりコマンドを実行できます。

コンテスト用クレートをworkspaceにまとめる

（追記）しばらくこれを使っていて、いろいろ不便なことが出てきたので、使うのをやめました。(たとえば、cargo checkやビルドの時間が長くなる、rust-analyzerがworkspace配下の他所のクレートのエラーを大量に拾って固まる、マラソンのテスターが上手く動かない、などなど…)
バイナリをどこかのディレクトリにまとめて容量を減らしたいだけであれば、後述の.cargo/config.tomlを使う方法が良いです。

一応、元の文章は残しておきます。

cargoにはworkspaceという機能があって、いくつかのクレートを1つの仮想的なクレート（？）の配下にまとめておくことができます。

workspace配下のクレートはバイナリを1か所に吐き出させることができるので、これでバイナリの容量を削減できます。
別にそんなことしなくて良くない？って思ってしまいそうですが、Rustはバイナリがデカいので、300回分のコンテストクレートを別々に作ってそれぞれビルドとかすると、数十GBとか下手すると100GB近い容量を消費することもザラです。（えぇ…）

workspaceを作るには、コンテスト用クレートをまとめるディレクトリを作り、そこでcargo initします。生成されたCargo.tomlを以下のような内容に書き換えれば完成です。

[workspace]
members = [
    "abc100", "abc101", ....(その他多数のクレート)
]
exclude = []
resolver = "2"

[workspace.dependencies]
ac-library-rs = "=0.1.1"
segtree = { git = "https://github.com/tayu0110/tayu-procon.git", package = "segtree" }
....(その他多数の依存クレート)

excludeとかresolverは明に指定しなくても動きます（確かresolverは指定しないとwarningが出た気がするが覚えていない…）。

workspace配下のクレートでは、今までのように依存クレートを宣言する以外に、以下のような宣言の仕方ができるようになります。

[dependencies]
ac-library-rs.workspace = true
segtree.workspace = true
.....

従来通りに指定した場合は依存クレートのバイナリがそのクレートのtarget/配下に配置され、workspaceのdependenciesを引き継ぐように指定した場合はworkspace内でバイナリが共有されます。このおかげで、各クレート配下に存在するバイナリの量を最低限にできます。

aとかbとか雑な名前つけてたけど、そんなバイナリをごったまぜにして一か所に吐かせて大丈夫なんか？っていう点ですが、cargoは賢いので、これを適切に処理してくれます。ただ、cargo-equipはこれを処理してくれないので、私は若干改造して使っています。

バイナリをまとめて容量を減らす

cargoの動作を制御するconfig.tomlをいうのを作ることで、バイナリを吐く宛先のディレクトリを変えることができます。
Configuration - The Cargo Book

コンテスト用クレート配下に.cargo/config.tomlを作成し、build.target-dirを書き込めばよいです。
私はコンテスト用クレートをまとめたatcoderというディレクトリの直下にtargetディレクトリを作り、そこにバイナリをまとめて吐かせたいので、

[build]
target-dir = "../target"

という設定ファイルを書き込んでいます。(Cargo.tomlからの相対パスなので、../targetとなります)

コンテスト用クレート作成スクリプト

ここまでで思ってる人多いと思いますが、めんどいですね。というわけでShellscriptで一発でコンテストごとのクレートを作れるようにしていきましょう。

私は以下のようなShellscriptを使っています。

#!/bin/bash -l

if [ -z "$1" ]; then
    echo "Usage: ./make_contest.sh [CONTEST_NAME]"
    exit
fi

DIR="$1"

if ls "$DIR" >/dev/null 2>&1; then
    echo "$DIR already exists."
    exit
fi

cargo new "$DIR"
cd "$DIR" || exit

mkdir .vscode -p
cat <<EOF >.vscode/settings.json
{
    "rust-analyzer.check.command": "check"
}
EOF

mkdir .cargo -p
cat <<EOF >.cargo/config.toml
[build]
target-dir = "../target"
EOF

cargo add ac-library-rs@=0.1.1
cargo add num@=0.4.1
cargo add rand@=0.8.5
cargo add regex@=1.9.1
cargo add permutohedron@=0.2.4
cargo add superslice@=1.0.0
cargo add itertools@=0.11.0
cargo add proconio@=0.4.5 --features derive

cargo add bitset --git https://github.com/tayu0110/tayu-procon.git
cargo add complex --git https://github.com/tayu0110/tayu-procon.git
cargo add convolution --git https://github.com/tayu0110/tayu-procon.git
cargo add ds --git https://github.com/tayu0110/tayu-procon.git --features "full"
cargo add fenwick-tree --git https://github.com/tayu0110/tayu-procon.git
cargo add flow --git https://github.com/tayu0110/tayu-procon.git
cargo add geometry --git https://github.com/tayu0110/tayu-procon.git
cargo add graph --git https://github.com/tayu0110/tayu-procon.git
cargo add math --git https://github.com/tayu0110/tayu-procon.git
cargo add matrix --git https://github.com/tayu0110/tayu-procon.git
cargo add mincost-flow --git https://github.com/tayu0110/tayu-procon.git
cargo add montgomery-modint --git https://github.com/tayu0110/tayu-procon.git
cargo add polynomial --git https://github.com/tayu0110/tayu-procon.git
cargo add rational --git https://github.com/tayu0110/tayu-procon.git
cargo add segtree --git https://github.com/tayu0110/tayu-procon.git
cargo add static-modint --git https://github.com/tayu0110/tayu-procon.git
cargo add string --git https://github.com/tayu0110/tayu-procon.git
cargo add suffix-array --git https://github.com/tayu0110/tayu-procon.git
cargo add unionfind --git https://github.com/tayu0110/tayu-procon.git
cargo add utility --git https://github.com/tayu0110/tayu-procon.git
cargo add wavelet-matrix --git https://github.com/tayu0110/tayu-procon.git

mkdir -p src/bin
for prefix in {a..g}; do
    SRC=src/bin/$prefix.rs
    {
        printf "use proconio::*;\n\n"
        printf "fn main() {\n"
        printf "    \n"
        printf "}\n"
    } >>"$SRC"
    cargo equip --exclude-atcoder-crates --exclude-atcoder-202301-crates --minify libs --no-rustfmt --no-check --remove docs --remove comments --bin "$prefix" >/dev/null
done

cargo build
cargo build --release

code .

.vscode/settings.jsonに打ち込んでいるのはRust Analyzerがリントを行うコマンドを指定するオプションです。
ライブラリや趣味プロではclippyという強力でアグレッシブなリントツールを使っていますが、コンテスト中だと鬱陶しいので、コンテスト用クレート内ではマイルドなcargo checkを使います。

最後の方ではcargo buildとcargo equipを回しています。1回目のビルドは外部クレートのビルドも走るので、めちゃくちゃ遅いです。
下手するとABCのA問題のコードを書き切るより時間がかかりますので、1回空回ししておくとストレスが減ります。