1 files changed, 131 insertions, 0 deletions
diff --git a/lib/segmenttree.ml b/lib/segmenttree.ml
new file mode 100644
index 00000000..2a7f9df0
--- /dev/null
+++ b/lib/segmenttree.ml
@@ -0,0 +1,131 @@
+(** This module is a very simple implementation of "segment trees".
+
+    A segment tree of type ['a t] represents a mapping from a union of
+    disjoint segments to some values of type 'a. 
+*)
+
+(** Misc. functions. *)
+let list_iteri f l =
+  let rec loop i = function
+    | [] -> ()
+    | x :: xs -> f i x; loop (i + 1) xs
+  in
+    loop 0 l
+
+let log2 x = log x /. log 2.
+
+let log2n x = int_of_float (ceil (log2 (float_of_int x))) 
+
+(** We focus on integers but this module can be generalized. *)
+type elt = int
+
+(** A value of type [domain] is interpreted differently given its position 
+    in the tree. On internal nodes, a domain represents the set of 
+    integers which are _not_ in the set of keys handled by the tree. On 
+    leaves, a domain represents the st of integers which are in the set of 
+    keys. *)
+type domain = 
+    (** On internal nodes, a domain [Interval (a, b)] represents 
+	the interval [a + 1; b - 1]. On leaves, it represents [a; b]. 
+	We always have [a] <= [b]. *)
+  | Interval of elt * elt
+      (** On internal node or root, a domain [Universe] represents all
+	  the integers. When the tree is not a trivial root,
+	  [Universe] has no interpretation on leaves. (The lookup
+	  function should never reach the leaves.) *)
+  | Universe
+
+(** We use an array to store the almost complete tree. This array
+    contains at least one element. *)
+type 'a t = (domain * 'a option) array
+
+(** The root is the first item of the array. *)
+let is_root i = (i = 0)
+
+(** Standard layout for left child. *)
+let left_child i = 2 * i + 1
+
+(** Standard layout for right child. *)
+let right_child i = 2 * i + 2
+
+(** Extract the annotation of a node, be it internal or a leaf. *)
+let value_of i t = match t.(i) with (_, Some x) -> x | _ -> raise Not_found
+
+(** Initialize the array to store [n] leaves. *)
+let create n init = 
+  Array.make (1 lsl (log2n n + 1) - 1) init
+
+(** Make a complete interval tree from a list of disjoint segments. 
+    Precondition : the segments must be sorted. *)
+let make segments = 
+  let nsegments = List.length segments in
+  let tree = create nsegments (Universe, None) in 
+  let leaves_offset = (1 lsl (log2n nsegments)) - 1 in
+
+    (** The algorithm proceeds in two steps using an intermediate tree
+	to store minimum and maximum of each subtree as annotation of
+	the node. *)
+
+    (** We start from leaves: the last level of the tree is initialized
+	with the given segments... *)
+    list_iteri 
+      (fun i ((start, stop), value) ->
+	 let k = leaves_offset + i in
+	 let i = Interval (start, stop) in
+	   tree.(k) <- (i, Some i))
+      segments;
+    (** ... the remaining leaves are initialized with neutral information. *)
+    for k = leaves_offset + nsegments to Array.length tree -1 do 
+      tree.(k) <- (Universe, Some Universe)
+    done;
+    
+    (** We traverse the tree bottom-up and compute the interval and
+	annotation associated to each node from the annotations of its
+	children. *)
+    for k = leaves_offset - 1 downto 0 do
+      let node, annotation = 
+	match value_of (left_child k) tree, value_of (right_child k) tree with
+	  | Interval (left_min, left_max), Interval (right_min, right_max) ->
+	      (Interval (left_max, right_min), Interval (left_min, right_max))
+	  | Interval (min, max), Universe -> 
+	      (Interval (max, max), Interval (min, max))
+	  | Universe, Universe -> Universe, Universe
+	  | Universe, _ -> assert false
+      in
+	tree.(k) <- (node, Some annotation)
+    done;
+
+    (** Finally, annotation are replaced with the image related to each leaf. *)
+    let final_tree = 
+      Array.mapi (fun i (segment, value) -> (segment, None)) tree
+    in
+      list_iteri 
+	(fun i ((start, stop), value) -> 
+	   final_tree.(leaves_offset + i) 
+	   <- (Interval (start, stop), Some value))
+	segments;      
+      final_tree
+
+(** [lookup k t] looks for an image for key [k] in the interval tree [t]. 
+    Raise [Not_found] if it fails. *)
+let lookup k t = 
+  let i = ref 0 in 
+    while (snd t.(!i) = None) do
+      match fst t.(!i) with
+	| Interval (start, stop) -> 
+	    if k <= start then i := left_child !i
+	    else if k >= stop then i:= right_child !i 
+	    else raise Not_found
+	| Universe -> raise Not_found
+    done;
+    match fst t.(!i) with
+      | Interval (start, stop) -> 
+	  if k >= start && k <= stop then
+	    match snd t.(!i) with 
+	      | Some v -> v
+	      | None -> assert false
+	  else 
+	    raise Not_found
+      | Universe -> assert false
+      
+