AtCoder ARC 120 C "Swap2"をRubyで
Rubyで解説している記事が一本くらいはあると楽観視していたのですが、なかったので私が書きます。何だってRubyistはこう中難易度問題に対して根性がねえんだ。
今回の問題はこれです。
アドホック部
「隣り合った左右の数字を入れ替えられる。ただし左から右に行った数字は-1、右から左に行った数字は+1される。最初の数列を最後の数列にするにはどうすれば良いか」
この問題でややこしい部分は右に行ったり左に行ったりする事で数字が増えたり減ったりする事です。これさえなければ、どの数字をどこにうごかせばいいのかひと目で判断できます。
どうすればいいでしょう?
ヒントは「2番目にある5と4番目にある3、どちらも一番左に持っていけば6。左から6番目に持っていけばどちらも1」です。
簡単ですね、与えられた数列にインデックス番号を加算します。配列A[3, 1, 4]と配列B[6, 2, 0]はそれぞれA'[3, 2, 6]と[6, 3, 2]となり、A[8, 5, 4, 7, 4, 5]とB[10, 5, 6, 7, 4, 1]はそれぞれA'[8, 6, 6, 10, 8, 10]とB'[10, 6, 8, 10, 8, 6]になります。どの数字をどこに持っていけばいいか目星をつけやすくなりましたね。
「いいや、まだ不十分だ。同じ数字が複数ある場合どれをどこに移動すればいいかわからない」という意見よーくわかります。そこで開始側の配列(配列A')をハッシュA''に圧縮します。配列の値をハッシュのキーに、配列のインデックスをハッシュの値の配列にプッシュします。意味不明だと言われてしまいそうですが、前述のA'[8, 6, 6, 10, 8, 10]をハッシュA''に圧縮すると{6 => [1, 2], 8 =>[0, 4], 10 => [3, 5]}となります。後は終了側の配列(配列B')を左から貪欲法で見ていき、ハッシュAの中で合致するインデックスのうち最も左のものを見ます。A側で10があるインデックスは3または5、このうちどちらかを一番左に持って来たい場合、普通は三番目を持ってきた方が操作数は少ないですよね。わざわざ五番目の10を持ってくる必要がありません。
こう考えると、配列A'[3, 2, 6]の数字はそれぞれ配列B'[6, 3, 2]のインデックス[1, 2, 0]に、A'[8, 6, 6, 10, 8, 10]の数字はそれぞれB'[10, 6, 8, 10, 8, 6]のインデックス[2, 1, 5, 0, 4, 3]に対応しています。
あとはこの[1, 2, 0]や[2, 1, 5, 0, 4, 3]をソートするには何回隣の数字を入れ替えればいいか……という話になります。
転倒数
[1, 2, 0]なら(1, 0)(2. 0)で転倒数は2。
[2, 1, 5, 0, 4, 3]なら(2, 1)(2, 0)(1, 0)(5, 0)(5, 4)(5, 3)(4, 3)で転倒数は7。
で、この転倒数は「隣の数字と入れ替えてソートする、つまりバブルソートするには何回の操作が必要か」の数字と一致するそうです。証明は知らん。(疲労困憊)
この転倒数はどう求めるかについては「2つの数字の左側の数字が大ループ、右の数字が小ループ」と「左は小ループで右が大ループ」の二択があります。どちらにしろ愚直法では二重ループになってしまって間に合いません。
では、どうするか。「位置の情報を捨てる」というのは思いつきそうです。[2, 1, 5, 0, 4, 3]で、「2つの数字の右が4」まで進んでいた場合、4より左の数字が2150と並んでいようが0125と並んでいようが関係ないわけですね。出現回数を数える配列を用意しておき、[1, 1, 1, 0, 0, 1]としておきインデックス4までの合計を取れば良いわけです……が、これでは二重ループを二重ループに書き換えただけで何の解決にもなっていません。
配列の部分和を高速で算出する方法は? 普通に考えれば累積和ですが、あれは前処理が重いですね。今回は元の数列が刻々と変わります。[1, 1, 1, 0, 0, 1]というのも右の数字が4の時の出現数配列であって、次の右の数が3の場合は出現数が[1, 1, 1, 0, 1, 1]となるわけです。いちいち累積和の配列を作り直していては本末転倒ですね。
元配列が変更された時の前処理が軽く、部分和を高速で算出できるデータ構造。そんな都合のいいものがあるかと言いたくなりますが、実はあります。フェニック木です。
なんか書いてるうちにセグ木で殴ってもよさそうな気がしてきましたが、勉強にならないのでフェニック木で解いていきます。
フェニック木(BIT木)
完全に1からの累積和を取っておくのではなく、程々に部分和を作っておく考え方です。そのときに2の累乗を目印にして部分和を作ると、元配列の加算にも部分和の計算にも計算量がO(logn)で均等に割り振られ、どちらか一方に負担がかからなくなります。
ここで筆者が体力の限界を迎えたため、元数列の項の一つに数字を足した場合の前処理であるaddと部分和を出すsum関数のアルゴリズムの説明は割愛します。参考リンクを読んでください。
それにしても、一個前の記事でやったあれってLSB(Least Significant Bit)って名前だそうですね。知りませんでした。
ACコード
14行目のループでAの値にインデックスを足してnew_a_valueとし、「この値はどことどことどこのインデックスで使われています」という情報をa_index_added_hashにメモしています。
22行目のループではBの値にインデックスを足してnew_b_valueとし(一個上のループでnew_a_valueと一緒にやれよ)、「その値はAで言うとどの位置にあったか」の配列をa_indexesに入れ、その条件が合致した場合のみ「その位置情報の中で最も左の数字を持ってくるのが一番操作数少ないな」となります。(B配列は左から見ていっているため、A配列で残っている中で最も左の物が最適になる。Aで右だった数字をBで左に、Aで左だった数字をBで右に持っていっても操作回数が増えるだけ。)合致しなかった場合は不可で終了。
「どこの数字をどこに移動させるか」のtarget_indexes配列が完成したら、出現回数をメモするtally_fenwick_tree配列を作ってフェニック木のモジュールを移譲します。
あとは最後のループ。フェニック木の特性上「n未満の数字の出現数」は高速で出せるので、全体 - n未満の数字の出現数をbuffer_answerに加算していって終了です。(数字の出現個数が1増えるので、フェニック木は更新を忘れないように)
Discussion