pythonで最短経路問題 with altgraph

pypiを眺めていて、altgraphというネットワークグラフ用のライブラリを見つけました。グラフ構造を割とシンプルに扱えるライブラリで、グラフ構造に付随するアルゴリズムも実装されています。早速使ってみようという事で、ダイクストラ法とベルマンフォード法のコードを書いてみました。

ライブラリの機能は殆ど使わず、グラフ構造の作成用途で使用しています。wikipediaの説明中にあったダイクストラ法の疑似コードが解り易かったので、同じスタイルで両者アルゴリズムを実装したものが以下になります。個人的にはグラフ構造を作成するのが面倒で仕方が無かったので、これからネットワークグラフが必要な時はこれを使おうと思います。

#!/usr/bin/python
# -*- coding: utf-8 -*-
 
import sys
from altgraph import Graph, Dot
 
def dijkstra(graph, start):
  dist = [sys.maxint for v in graph.nodes]
  prev = [None for v in graph.nodes]
  dist[start] = 0
  Q = graph.nodes.copy()
 
  while len(Q) > 0:
    (u,minimum) = min([(v,dist[v]) for v in Q], key=lambda x:x[1])
    neighbors = [graph.edges[e] for e in Q[u][0]]
    del(Q[u])
    for v,tmp,distance in neighbors:
      alt = dist[u] + distance
      if alt < dist[v]:
        (dist[v], prev[v]) = (alt, u)
 
  return dist, prev
 
def bellmanford(graph, start):
  dist = [sys.maxint for e in graph.nodes]
  prev = [None for e in graph.nodes]
  dist[start] = 0
 
  for v in graph.nodes:
    for e in graph.edges:
      (u, v, weight) = graph.edges[e]
      if dist[u] + weight < dist[v]:
        (dist[v], prev[v]) = (dist[u] + weight, u)
 
  for e in graph.edges:
    (u, v, weight) = graph.edges[e]
    if dist[u] + weight < dist[v]:
      raise "Graph contains a negative-weight cycle"
 
  return dist, prev
 
if __name__ == '__main__':
  # グラフの定義
  edges = [(0,1,2), (0,2,4), (0,3,5),  # (from, to, length)
           (1,2,3), (2,3,2), (1,4,6),
           (2,4,2), (3,5,6), (4,5,4)]
  # グラフ表現の作成  
  graph = Graph.Graph()
  for (frm, to, length) in edges:
    graph.add_edge(frm, to, length)
    graph.add_edge(to, frm, length)
 
  # 最短距離を求める
  start  = 0
  target = 5
  print bellmanford(graph, start)[0][target]
  print dijkstra(graph, start)[0][target]
 
  # グラフを描画する
  dot = Dot.Dot(graph)
  dot.display()