AtCoder ABC049 D (ARC065) 連結 Union-Find木

　ABC044の山を越えるとあとはなんというか手応えのない問題ばかりでここにまとめるほどのこともなかったのだが、50に来る前にちょっと骨のあるのが来た。骨といっても典型問題のようで、アルゴリズムを覚えれば一発らしいがこっちは無知(無恥？)なので Python3 ではなかなかつらかった。ということで記録を残しておく。

問題
経路探索？
集合
Union-Find木

問題

　都市がN個ある。それぞれをつなぐ鉄道網と道路網を与えるのでそれぞれの都市から鉄道と道路で行ける都市の数を数えろ。最小はその都市自体の1。

経路探索？

　お？これは ABC014 D 閉路閉路で出てきた経路問題か？あれは結局 Python3 でACできず、とりあえず PyPy でごまかしてきたのだ。つらい。つらいが目の前に来た問題はそれなりになんとかこなさねばならない。ということでとりかかる。とりあえずDFSっぽい探索で毎度いちいちどこまで行けるのか探索するようなものを書いてみたが当然ながらTLE。遅い Python3 でいちいち重いループを書いたら即死なのだ。(というか普通は計算量Oで見積る) うーむむ。

集合

　というところで気がついた。道路でつながってるところは全部の都市それぞれに行けるのだからつながってる都市間で行ける到達可能都市数は同じではないか。鉄道もおなじ。ということは網でつながってるのは全部別々のかたまりで表現できるのではないか。道路でも鉄道でも行ける都市ということはその集合同士が重複してるところではないか。つまり道路集合かつ鉄道集合なのではないか。

f:id:inudaisho:20181018184514p:plain

　よしよし。グループ分けすればいいんだな。ということで、グループに番号をつけてみたがやっぱりTLE。うーむ。解法をみてもどうもつまり集合だろ？ということしか書いてない。

　そもそも Python3 は遅いのでいちいち各都市から行ける都市の辞書をつくって、そこから交通網別の都市の集合をつくって、都市からグループを引ける辞書つくって、とやってるとACできないのだ。順番に処理しながらなにか構造をつくってそこから一発で必要なものを引っぱるような仕組みをつくらないとうまくいかない。Python3 でACしてる人もたいてい1000ms以上かかっている。これはなにかアルゴリズムがあるな、と人の解答をみてみると UnionFind とかいうクラスつくったりfindとかunionとかいう関数つくったりしてる人が多い。C++ の方をみると速い人は18msで出してるが、中国人なのにコメントが日本語で書いてあり、コメント部分だけみると次点の23msの人とおなじなので、改良したかコピペして新しくなった鯖の性能で稼いだかどっちかだろう。自分も平気で人の成果をパク...応用するので人のことは言えない。勉強勉強。さてどんなアルゴリズムだろう。

Union-Find木

　ということでググるとすぐに出てきた。Union-Find 木。最初書いた集合のように互いに重ならない集合のことを素集合といい、それを表現するのによく使われるらしいが非常に有用であるようだ。

素集合データ構造 - Wikipedia

　素集合からわかりやすく説明してるのがここですねー。Pythonやってる人の説明はざっくりしたところからそこそこ細かいところまで押えてるのが多くて自分のような物好きにはありがたい。

at274.hatenablog.com

　あとAtCoder公式の技術寄りの説明がこれ。公式がクドクド説明しだしたら終わりなのでこれでいいとおもいます。

Union find(素集合データ構造) from AtCoder Inc.

www.slideshare.net

　この問題の場合、各交通網の代表都市を決めて、そこにつながるような木を構築するということらしい。根(代表都市)を探すのがfind、別々の木を結合するのがunionで、あわせてunionfindと総称されるそうだ。特に高速にするには、各都市からその代表都市を参照できるように網を張りかえるとよいらしい。えーとでもそれって結合するときに全部張りかえるということ？却って遅くないか？

　とおもうが、集合のところに書いたように、いちいちこれをやってあれをやってと段階毎にキレイにやろうとするから遅くなるのだ。図書館に行くついでにスーパーで野菜を買えば図書館とスーパー別々に行くより手間がはぶける。ということで、A集団とB集団を結合してB集団の方を根にするには、A集団の代表都市(根)からB集団の代表都市(根)へリンクをはってA集団をB集団の枝にしておき、リンクの張り替えはfindのときついでにやるようにする。木構築が中途半端で終わるので、この段階では各都市がすべて代表都市につながってるとは言い難い状態だが、この場合、構築した木を利用してより複雑なことをするのではなく、各都市がどの集合に属しているか調べるだけ。つまりどうせ最後に全都市について代表都市(根)を検索するのでその最中に張り替えが完成する。という具合だ。

　ということでその方向で実装したのがこちら。だいたいみんな最小の方に代表都市を置いてるようなので最大の方に代表都市を置いてみた。たいしてかわらんだろ。

beta.atcoder.jp

　700ms台でました!