AtCoder Beginner Contest 016 C. 友達の友達:解法とPython 3による解答例
最近、私はAtCoder Beginner Contest (ABC)およびAtCoder Regular Contest (ARC)の過去問を解き、競技プログラミングの練習を行っています。その一環として解いたC: 友達の友達 - AtCoder Beginner Contest 016 | AtCoderという問題について、この記事ではその解法および解答例を、学習の記録を兼ねてまとめておきます。
前提:最短経路問題とそれを解くアルゴリズム*1
解法について述べる前に、それが基づいている最短経路問題とそれを解くためのアルゴリズムについて簡単にまとめておきます。なお、アルゴリズムの詳細については、以下に示すそれを実装したコード中のコメントで適宜補足しています。
ある2頂点が与えられたときに、それらを始終点とする経路であって、通る辺の重みの和が最小となるものを、その2頂点間の最短経路(shortest path)と呼び、これを求める最適化問題を最短経路問題といいます*2。最短経路問題の代表例としては、次の2つが挙げられます*3。
- 全点対最短経路問題(all-pairs shortest path problem)
- 単一始点最短経路問題(single-source shortest path problem)
全点対最短経路問題とWarshall-Floyd法*4
全点対最短経路問題は、グラフ内のあらゆる2頂点間の最短経路を求める問題であり、
解法
以上を念頭に置くと、C: 友達の友達 - AtCoder Beginner Contest 016 | AtCoderは次のような方法で解けます。SNSのユーザを、グラフの頂点とみなせば、ユーザiとユーザjが友達同士であるということは、頂点iから頂点jに、また頂点jから頂点iに、それぞれ重み1の辺が存在するということと同一視できます。すると、ユーザiとjが「友達の友達」であることは、頂点iから頂点jまで(その逆も同様)の最短経路の長さが2であるということになります。そこで、Warshall-Floyd法により全点対最短経路を求めておき、各頂点に対してそれとの最短経路長が2である頂点を数えればよいといえます。あるいは、Bellman-Ford法やDijkstra法を各頂点に対して適用し、同様の数え上げを行っても答えが得られます。どちらかといえばWarshall-Floyd法を使う方が簡便でしょう。なお、本問ではとみなせます。
解答例
以上の解法に基づく、Python 3による解答コードを示しておきます。
Warshall-Floyd法による解答例
以下にWarshall-Floyd法を用いたコードを示します。計算量はとなります。
Bellman-Ford法による解答例
次はBellman-Ford法を繰り返し用いたものです。計算量はとなります。
別解
上では最短経路問題を解くための代表的アルゴリズムを利用して答えを求めましたが、そのようなアルゴリズムを利用せず、どちらかといえばad hocな手法によることもできます。公式の解説スライドAtCoder Beginner Contest 016 解説にその例が記されています。
参考にした書籍・ウェブページ
Warshall-Floyd法、Bellman-Ford法およびDijkstra法を利用した上記のコードは、以下の図書を参考にして作成したものです。
- 『プログラミングコンテストチャレンジブック』(第2版)2.5節
- 通称:蟻本
- 『アルゴリズムクイックリファレンス』6.4節および6.5節
また、Warshall-Floyd法の実装にあたっては、次の記事がとても参考になりました。
pakapa104.hatenablog.com
感想
少し前までは、グラフはなんだか難しそうだと思い込み敬遠しがちでした。しかしながら、最近では問題演習を通じて、グラフやそれにまつわるアルゴリズムが少しずつ身近になってきました。今回、最短経路問題に関するアルゴリズムを学んだことで、より一層グラフが「こわく」なくなったとともに、解けそうな問題が増えたように感じました。今後ともコツコツと問題を解き、「実装力」を鍛えるとともに、さまざまな技法を学んでいきたいところです。
*1:グラフ理論に関する基本的用語については、wikipedia:グラフ理論をはじめとするウェブページや、プログラミングコンテストチャレンジブック(第2版)などの成書を参照してください。
*2:wikipedia:最短経路問題および『プログラミングコンテストチャレンジブック』(第2版)p.94
*3:wikipedia:最短経路問題およびwikipedia:en:Shortest_path_problem
*4:前掲書p.98
*5:前掲書pp.95-97
*6:wikipedia:ダイクストラ法によります。同記事によると、優先度付きキューの代わりに単なる集合を用いる場合には、計算量はとなるといいます。