重み付き木構造上の2点間の距離を求めるアルゴリズムを python で書いてみる

グラフの一種である木構造での距離を求めるアルゴリズムを考えます.

グラフとは

詳しくはグラフ理論 - Wikipedia などの記述をご参照ください.
グラフ(graph)とは頂点(vertex)と辺(edge)からなる構造を表しており, 頂点同士が辺によって繋がっていることを表しているものです（頂点の集合は自然数の部分集合と同型なので, それぞれの頂点には何番目の頂点かを表す自然数が割り当てられます. 一般には辺それぞれには実数の重みがあり, その頂点同士の移動にかかるコストを表しています. 辺は $(v_1, v_2, c)$ の形式を持ち, 頂点 $v_1$ から頂点 $v_2$ はコスト $c$ で隣接していることを表しています. 重み付きでない場合は, $c=1$ と理解します.
辺に向きがあるもの（つまり, 2つの頂点について行けるのに戻れなかったりコストが等しくなかったりするもの）を有向グラフ(directed graph), 辺に向きがないもの（つまり, 頂点間のコストは行きと帰りで等しいもの）を無向グラフ(undirected graph)といいます. 無向グラフでは $(v_1, v_2, c)$ なる辺があったとき, 逆向きの $(v_2, v_1, c)$ を辺にもつ有向グラフと解釈することもできます.

木とは

まずことばの定義をします. 無向グラフにおいて, 始点と終点を表す2つの頂点に対して, 始点から終点に辺と頂点を伝って行けるとき, その道順をパス(path)といいます(正確には, 隣接する頂点の列を言います). グラフ上の任意の2つの頂点に対して, その2点を結ぶパスがあるとき, グラフは連結(connected)であるといいます. また, 始点と終点が一致するパスを閉路(cycle)といいます.
木とは, 無向グラフで連結かつ閉路を持たないものとして定義します.

wikipedia などの記述を見ても分かるとおり, グラフには色々な種類（とことば）があります. （今はやりの？）ニューラルネットもグラフとして記述されることが多く, 非常に重要な対象です. 徐々に解説できればと思います.

スタックについて

リストやセット, 連想配列のように, コンピュータの中にデータを系統的に格納する（つまり, メモリやハードディスクに情報をどのように乗せるかという）形式をデータ構造といいます. 実は上で紹介したグラフもデータ構造です. データ構造とその上でのアルゴリズムを考えることによってプログラムのパフォーマンスが大きく変わることがあります.
今日紹介するスタック(stack)はその中でも代表的なもので, データの挿入と取り出しに特化したリストのようなものです. 言い換えると, "push" で要素をリストの先頭に挿入し, "pop" でリストの先頭にある要素を取得する配列です. つまり, 最後に挿入した要素が常に取得されます. データの後入れ先出し(Last In First Out)がされるということで, この構造をLIFOと呼びます.

以上の準備の下, 重み付き木構造が与えられたとき, 始点からそれぞれの頂点への距離を求めるアルゴリズムを考えていきます.
まず始点 $v_0$ , 辺の集合 $\textrm{edges}$ , 頂点の個数 $N$ が与えられたとき, 終点 $w$ までの距離を次のフローで求めます.
1. $\textrm{dist}(k)\text{ for } k=1,\cdots, N$ を始点 $v_0$ から頂点 $k$ までの距離を表す. これを求めたい.
2. 末端点を表すスタックに $v_0$ を入れ, $\textrm{dist}(v_0)=0$ とする.
3. スタックから要素 $v_1$ を取り出し, $v_1$ を含み, なおかつ $v_2$ が未訪問であるような辺 $(v_1, v_2, c)$ について
$\textrm{dist}(v_2)=\textrm{dist}(v_1)+c$
とする.
4. スタックに $v_2$ を新たに挿入する.
5. 3～4を繰り返し, $\textrm{dist}(w)$ を求める.
となります. pythonコードは

def BFS(K,edges,N):
  roots=[ [] for i in range(N)]
  for a,b,c in edges:
    roots[a-1]+=[(b-1,c)]
    roots[b-1]+=[(a-1,c)]
  dist=[-1]*N
  stack=[]
  stack.append(K)
  dist[K]=0
  while stack:
    label=stack.pop(-1)
    for i,c in roots[label]:
      if dist[i]==-1:
        dist[i]=dist[label]+c
        stack+=[i]
      if i==w:
        return dist[w]

となろうかと思います. roots は頂点 $i$ から頂点 $j$ にコスト $c$ て行く辺を集めているので, $\textrm{roots}(i)$ は [(j,c)] を辺によって隣接する要素とする集合としています. したがって roots は配列の配列になっています. また, 初期化 dist=[-1]*N によって未訪問の頂点のフラグを-1としています. 関数を BFS としたのは幅優先探索(Breedth First Searc)によって最短距離を求めているからです. スタックを python の標準で与えられているリストで実装していますが, append と pop(-1) を組み合わせることでスタックとして扱うことが出来ます. しかも, それぞれの計算量はリストの大きさに依存しません. したがって今回のコードの計算量は, while 文をまわすのにかかる $O(N)$ です.
今回は関係ありませんが, キュー(que)というデータの先入れ先出しを行うデータ構造を用いたい場合には, リストでは要素の取り出しまたは追加にリストの大きさに比例する時間がかかってしまいます.

前回と同様, 次回の記事で競技プログラミングの問題を考えていきます.