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頂点間の最短経路を求める問題であり、Warshall-Floyd（ワーシャルフロイド）法は、これを解くための動的計画法に基づくアルゴリズムです。頂点全体の集合をとし、その要素数（すなわち頂点の数）をと表すと、この方法の計算量はとなります。Warshall-Floyd法は、重みが負の辺があっても動作します。また、このアルゴリズムで負の閉路、つまり辺の重みの総和が負となる閉路を検出することもできます。

単一始点最短経路問題とBellman-Ford法およびDijkstra法*5

単一始点最短経路問題は、始点を固定したときに、それとは異なるすべての頂点との間の最短経路を求める問題です。これを解くための代表的アルゴリズムとしてはBellman-Ford（ベルマンフォード）法とDijkstra（ダイクストラ）法が挙げられます。両者の特徴を以下にまとめておきます。なお、辺全体の集合をとし、その要素数（すなわち辺の数）をと表しています。

• Bellman-Ford法は重みが負である辺が含まれているグラフに適用可能であるが、Dijkstra法は適用不可能である
• Bellman-Ford法は負の閉路の検出に使用可能である
• Bellman-Ford法の計算量はであるが、優先度付きキュー（二分ヒープ）を用いたDijkstra法のそれはである*6

解法

以上を念頭に置くと、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法を繰り返し用いたものです。計算量はとなります。

Dijkstra法による解答例

最後に、優先度付きキューを用いたDijkstra法を繰り返し用いて計算するコードを示します。Python標準ライブラリのheapqを用いれば、リストを優先度付きキュー（二分ヒープ）として手軽に使うことができます。計算量はとなります。

別解

上では最短経路問題を解くための代表的アルゴリズムを利用して答えを求めましたが、そのようなアルゴリズムを利用せず、どちらかといえばad hocな手法によることもできます。公式の解説スライドAtCoder Beginner Contest 016 解説にその例が記されています。

参考にした書籍・ウェブページ

Warshall-Floyd法、Bellman-Ford法およびDijkstra法を利用した上記のコードは、以下の図書を参考にして作成したものです。

• 『プログラミングコンテストチャレンジブック』（第2版）2.5節
  • 通称：蟻本
• 『アルゴリズムクイックリファレンス』6.4節および6.5節
  • 表紙にヤドカリが描かれている本です。これらの節では、上記の3つのアルゴリズムとその応用について、C++による実装例とともに詳述されています。

また、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:ダイクストラ法によります。同記事によると、優先度付きキューの代わりに単なる集合を用いる場合には、計算量はとなるといいます。