OCamlでダイクストラ法

30分プログラム、その692。OCamlダイクストラ法。
"OCaml ダイクストラ法"でググると、昔ボクが書いた不完全な実装がヒットしてしまう。(ダイクストラ法を実装しようとしたら、よくわからないものになった - みずぴー日記)
この不完全な実装を放置するのはよくない気がしたので、ちゃんと実装し直しました。
前回は純粋関数的に書こうとして失敗したので、今回はmutableなフィールドを持っているレコードを使いました。ただし、モジュール内部で隠蔽して、外部からだとまるで副作用が無いかのように扱えるようにしました。

シグネチャ

type 'a graph
type 'a node
type 'a edge

(* グラフの構築 *)
val make_node : 'a -> 'a node
val make_edge : 'a node -> 'a node -> int -> 'a edge
val make_graph : nodes:'a node list -> edges:'a edge list -> 'a graph

(* ノードには任意のデータを持たせれるので、それを取り出す *)
val node_data : 'a node -> 'a

(* 最短距離と経路を求める *)
val shortest : 'a graph -> 'a node -> 'a node -> (int * 'a node list) option

使い方

(* ノードを作る *)
let a = make_node "A"
let b = make_node "B"
let c = make_node "C"
let d = make_node "D"
let e = make_node "E"

(* グラフを構築する *)
let nodes = [
  a; b; c; d; e
]
let edges = [
  make_edge a b 3;
  make_edge b c 1;
  make_edge a c 5;
  make_edge c d 1;
]

let graph = make_graph ~edges ~nodes

(* aからdの最短距離と経路を求める *)
let (distance,path) = shortest a d

実装

(*
compile:

  ocamlfind ocamlc -package extlib -linkpkg dijkstra.mli dijkstra.ml

example:

  let a = make_node "A"
  let b = make_node "B"
  let c = make_node "C"
  let d = make_node "D"
  let e = make_node "E"

  let nodes = [
    a; b; c; d; e
  ]
  let edges = [
    make_edge a b 3;
    make_edge b c 1;
    make_edge a c 5;
    make_edge c d 1;
  ]

  let graph = make_graph ~edges ~nodes

  let (distance,path) = shortest a d
  (* where distance=5 path=[a;b;c;d] *)
*)

open StdLabels

let (@@) f g = f g
let (+>) f g = g f

let sure f =
  function
      Some x ->
	Some (f x)
    | None ->
	None

type 'a node = {
  data: 'a;
  mutable path : int * 'a node list option
}

type 'a edge = {
  from: 'a node;
  to_ : 'a node;
  distance : int
}

type 'a graph = {
  nodes : 'a node list;
  edges : 'a edge list;
  mutable start_node : 'a node option
}

let make_node data = {
  data = data;
  path = (max_int,None)
}

let make_edge from to_ distance = {
  from=from;
  to_=to_;
  distance=distance
}

let node_data {data=data} = data

let make_graph ~nodes ~edges = {
  start_node = None;
  nodes = nodes;
  edges = edges
}

let edges {edges=edges} = edges
let nodes {nodes=nodes} = nodes

let min_node a b =
  if fst a.path < fst b.path then a else b

let minimum_node nodes =
  List.fold_left ~f:min_node ~init:(List.hd nodes) (List.tl nodes)

let connect_edges {edges=edges} node =
  List.filter ~f:(fun {from=from} -> node = from) edges

let cons x xs = x :: xs

let rec update_nodes graph nodes =
  if nodes = [] then
    ()
  else
    let { path=(d, path) } as node =
      minimum_node nodes in
    let edges =
      connect_edges graph node in
      List.iter edges ~f:begin fun {distance=distance; to_=to_} ->
	if distance + d < fst to_.path  then
	  to_.path <-
	    (distance + d, sure (cons node) path)
      end;
      update_nodes graph (ExtList.List.remove nodes node)

let shortest graph first last =
  if graph.start_node <> Some first then begin
    first.path <- (0,Some []);
    update_nodes graph graph.nodes;
    graph.start_node <- Some first
  end;
  match last.path with
      distance,Some path ->
	Some (distance,List.rev (last::path))
    | _,None ->
	None