dijkstra.ml 2.26 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
(* This module implements Dijkstra's algorithm over a graph with labeled
   weighted edges. (We assume the weight is contained in the label.) *)

module Make (G : sig

  (* This is the type of graph vertices. *)

  type vertex
  include Hashtbl.HashedType with type t := vertex

  (* This is the type of graph labels. *)

  type label

15
  (* The source vertex (or vertices). *)
16

17
  val sources: (vertex -> unit) -> unit
18

19
  (* The weighted outgoing edges of a vertex. *)
20

21
  val successors: vertex -> (label -> int -> vertex -> unit) -> unit
22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53

end) = struct
  open G

  (* The priority queue contains triples of a total weight, a vertex,
     and the discovery path for this vertex -- a list of labels. *)

  module Element = struct
    type t = int * vertex * label list
    let compare (w1, _, _) (w2, _, _) =
      Pervasives.compare w2 w1 (* switched because [Heap] uses the wrong convention *)
  end

  module PQ =
    Heap.Imperative(Element)

  (* A hash table maps vertices to distances. Another table maps vertices to unit
     values, indicating whether this distance is final. *)

  module H =
    Hashtbl.Make(struct include G type t = vertex end)

  let search f =

    (* Create the data structures. *)
    let queue = PQ.create 1024 in
    let dist : int H.t = H.create 1023 in
    let fixed : unit H.t = H.create 1023 in
    (* A handy accessor. *)
    let distance v =
      try H.find dist v with Not_found -> max_int
    in
54 55 56 57 58
    (* Insert the initial vertices. *)
    sources (fun source ->
      PQ.add queue (0, source, []);
      H.add dist source 0
    );
59 60 61 62 63 64 65 66 67 68 69
    (* As long as the queue is nonempty, ... *)
    while not (PQ.is_empty queue) do
      (* Extract one vertex. *)
      let (w, v, p) = PQ.pop_maximum queue in
      (* Check if this vertex is final already. If so, nothing to do. *)
      if not (H.mem fixed v) then begin
        (* Mark it final. *)
        H.add fixed v ();
        (* Let the client know about it. *)
        f (w, v, p);
        (* Examine its outgoing edges. *)
70
        successors v (fun label weight v' ->
71
          let w' = weight + w in
72 73 74 75 76
          if w' < distance v' then begin
            assert (not (H.mem fixed v'));
            H.replace dist v' w';
            PQ.add queue (w', v', label :: p)
          end
77
        )
78
      end
79 80
    done;
    distance
81 82 83

end