functional SplayTree as Maps

splayMap.ml

@@ -0,0 +1,198 @@
+(*
+Copyright (c) 2013, Simon Cruanes
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.  Redistributions in binary
+form must reproduce the above copyright notice, this list of conditions and the
+following disclaimer in the documentation and/or other materials provided with
+the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*)
+
+(** {1 Functional Maps} *)
+
+(* We use splay trees, following
+http://www.cs.cornell.edu/Courses/cs3110/2009fa/recitations/rec-splay.html
+*)
+
+
+type ('a, 'b) t = {
+  cmp : 'a -> 'a -> int;
+  mutable tree : ('a, 'b) tree;  (* for lookups *)
+} (** Tree with keys of type 'a, and values of type 'b *)
+and ('a, 'b) tree =
+  | Empty
+  | Node of ('a * 'b * ('a, 'b) tree * ('a, 'b) tree)
+
+let empty_with ~cmp =
+  { cmp;
+    tree = Empty;
+  }
+
+let empty () =
+  { cmp = Pervasives.compare;
+    tree = Empty;
+  }
+
+let is_empty t =
+  match t.tree with
+  | Empty -> true
+  | Node _ -> false
+
+(** Pivot the tree so that the node that has key [key], or close to [key], is
+    the root node. *)
+let rec splay ~cmp (k, v, l, r) key =
+  let c = cmp key k in
+  if c = 0
+    then (k, v, l, r) (* found *)
+  else if c < 0
+    then match l with
+    | Empty -> (k, v, l, r)  (* not found *)
+    | Node (lk, lv, ll, lr) ->
+      let lc = cmp key lk in
+      if lc = 0
+        then (lk, lv, ll, Node (k, v, lr, r))  (* zig *)
+      else if lc < 0
+        then match ll with
+        | Empty -> (lk, lv, Empty, Node (k, v, lr, r)) (* not found *)
+        | Node n -> (* zig zig *)
+          let (llk, llv, lll, llr) = splay ~cmp n key in
+          (llk, llv, lll, Node (lk, lv, llr, Node (k, v, lr, r)))
+      else
+        match lr with
+        | Empty -> (lk, lv, ll, Node (k, v, Empty, r))
+        | Node n -> (* zig zag *)
+          let (lrk, lrv, lrl, lrr) = splay ~cmp n key in
+          (lrk, lrv, Node (lk, lv, ll, lrl), Node (k, v, lrr, r))
+  else match r with
+    | Empty -> (k, v, l, r)  (* not found *)
+    | Node (rk, rv, rl, rr) ->
+      let rc = cmp key rk in
+      if rc = 0
+        then (rk, rv, Node (k, v, l, rl), rr)  (* zag *)
+      else if rc > 0
+        then match rr with
+        | Empty -> (rk, rv, Node (k, v, l, rl), Empty)  (* not found *)
+        | Node n ->  (* zag zag *)
+          let (rrk, rrv, rrl, rrr) = splay ~cmp n key in
+          (rrk, rrv, Node (rk, rv, Node (k, v, l, rl), rrl), rrr)
+      else match rl with
+        | Empty -> (rk, rv, Node (k, v, l, Empty), rr) (* zag zig *)
+        | Node n ->  (* zag zig *)
+          let (rlk, rlv, rll, rlr) = splay ~cmp n key in
+          (rlk, rlv, Node (k, v, l, rll), Node (rk, rv, rlr, rr))
+
+let find t key =
+  match t.tree with
+  | Empty -> raise Not_found
+  | Node (k, v, l, r) ->
+    let (k, v, l, r) = splay ~cmp:t.cmp (k, v, l, r) key in
+    t.tree <- Node (k, v, l, r);   (* save balanced tree *)
+    if t.cmp key k = 0
+      then v
+      else raise Not_found
+
+let mem t key =
+  match t.tree with
+  | Empty -> false
+  | Node (k, v, l, r) ->
+    let (k, v, l, r) = splay ~cmp:t.cmp (k, v, l, r) key in
+    t.tree <- Node (k, v, l, r);   (* save balanced tree *)
+    if t.cmp key k = 0
+      then true
+      else false
+
+(** Recursive insertion of key->value in the tree *)
+let rec insert ~cmp tree key value =
+  match tree with
+  | Empty -> Node (key, value, Empty, Empty)
+  | Node (k, v, l, r) ->
+    let c = cmp key k in
+    if c = 0
+      then Node (key, value, l, r)  (* replace *)
+    else if c < 0
+      then Node (k, v, insert ~cmp l key value, r)
+    else Node (k, v, l, insert ~cmp r key value)
+
+let add t key value =
+  let tree =
+    match t.tree with
+    | Empty -> Node (key, value, Empty, Empty)
+    | Node (k, v, l, r) ->
+      let (k, v, l, r) = splay ~cmp:t.cmp (k, v, l, r) key in
+      let tree = Node (k, v, l, r) in
+      t.tree <- tree;  (* save balanced tree *)
+      (* insertion in this tree *)
+      insert ~cmp:t.cmp tree key value
+  in
+  { t with tree; }
+
+let singleton ~cmp key value =
+  add (empty_with ~cmp) key value
+
+(** Merge of trees, where a < b *)
+let rec left_merge a b =
+  match a, b with
+  | Empty, Empty -> Empty
+  | Node (k, v, l, r), b -> Node (k, v, l, left_merge r b)
+  | Empty, b -> b
+
+let remove t key =
+  match t.tree with
+  | Empty -> t
+  | Node (k, v, l, r) ->
+    let (k, v, l, r) = splay ~cmp:t.cmp (k, v, l, r) key in
+    t.tree <- Node (k, v, l, r);
+    if t.cmp key k = 0
+      then (* remove the node, by merging the subnodes *)
+        let tree = left_merge l r in
+        { t with tree; }
+      else (* not present, same tree *)
+        t
+
+let iter t f =
+  let rec iter t = match t with
+  | Empty -> ()
+  | Node (k, v, l, r) ->
+    f k v;
+    iter l;
+    iter r
+  in iter t.tree
+
+let fold t acc f =
+  let rec fold acc t = match t with
+  | Empty -> acc
+  | Node (k, v, l, r) ->
+    let acc = f acc k v in
+    let acc = fold acc l in
+    fold acc r
+  in
+  fold acc t.tree
+
+let size t = fold t 0 (fun acc _ _ -> acc+1)
+
+let choose t =
+  match t.tree with
+  | Empty -> raise Not_found
+  | Node (k, v, _, _) -> k, v
+
+let to_seq t =
+  Sequence.from_iter
+    (fun kont -> iter t (fun k v -> kont (k, v)))
+
+let of_seq t seq =
+  Sequence.fold (fun t (k, v) -> add t k v) t seq

splayMap.mli

@@ -0,0 +1,69 @@
+(*
+Copyright (c) 2013, Simon Cruanes
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+Redistributions of source code must retain the above copyright notice, this
+list of conditions and the following disclaimer.  Redistributions in binary
+form must reproduce the above copyright notice, this list of conditions and the
+following disclaimer in the documentation and/or other materials provided with
+the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*)
+
+(** {1 Functional Maps} *)
+
+type ('a, 'b) t
+  (** Tree with keys of type 'a, and values of type 'b *)
+
+val empty_with : cmp:('a -> 'a -> int) -> ('a, 'b) t
+  (** Empty tree *)
+
+val empty : unit -> ('a, 'b) t
+  (** Empty tree using Pervasives.compare *)
+
+val is_empty : (_, _) t -> bool
+  (** Is the tree empty? *)
+
+val find : ('a, 'b) t -> 'a -> 'b
+  (** Find the element for this key, or raises Not_found *)
+
+val mem : ('a, _) t -> 'a -> bool
+  (** Is the key member of the tree? *)
+
+val add : ('a, 'b) t -> 'a -> 'b -> ('a, 'b) t
+  (** Add the binding to the tree *)
+
+val singleton : cmp:('a -> 'a -> int) -> 'a -> 'b -> ('a, 'b) t
+  (** Singleton map *)
+
+val remove : ('a, 'b) t -> 'a -> ('a, 'b) t
+  (** Remove the binding for this key *)
+
+val iter : ('a, 'b) t -> ('a -> 'b -> unit) -> unit
+  (** Iterate on bindings *)
+
+val fold : ('a, 'b) t -> 'c -> ('c -> 'a -> 'b -> 'c) -> 'c
+  (** Fold on bindings *)
+
+val size : (_, _) t -> int
+  (** Number of bindings (linear) *)
+
+val choose : ('a, 'b) t -> ('a * 'b)
+  (** Some binding, or raises Not_found *)
+
+val to_seq : ('a, 'b) t -> ('a * 'b) Sequence.t
+
+val of_seq : ('a, 'b) t -> ('a * 'b) Sequence.t -> ('a, 'b) t

tests/run_tests.ml

@@ -6,6 +6,7 @@ let suite =
   "all_tests" >:::
     [ Test_pHashtbl.suite;
       Test_PersistentHashtbl.suite;
+      Test_splayMap.suite;
       Test_leftistheap.suite;
       Test_cc.suite;
       Test_puf.suite;

tests/test_splayMap.ml

@@ -0,0 +1,41 @@
+
+open OUnit
+
+let test1 () =
+  let empty = SplayMap.empty () in
+  let m = SplayMap.of_seq empty (Sequence.of_list [1, "1"; 2, "2"; 3, "3"]) in
+  OUnit.assert_equal ~printer:(fun s -> s) "2" (SplayMap.find m 2);
+  OUnit.assert_equal ~printer:(fun s -> s) "2" (SplayMap.find m 2);
+  OUnit.assert_equal ~printer:(fun s -> s) "3" (SplayMap.find m 3);
+  OUnit.assert_equal ~printer:(fun s -> s) "1" (SplayMap.find m 1);
+  OUnit.assert_raises Not_found (fun () -> SplayMap.find m 4);
+  ()
+
+let test_remove () =
+  let n = 100 in
+  let m = SplayMap.of_seq (SplayMap.empty ())
+    (Sequence.zip (Sequence.zip_i (Sequence.int_range ~start:0 ~stop:n))) in
+  for i = 0 to n do
+    OUnit.assert_equal ~printer:string_of_int i (SplayMap.find m i);
+  done;
+  let m = SplayMap.remove m (n/2) in
+  OUnit.assert_equal ~printer:string_of_int n (SplayMap.find m n);
+  OUnit.assert_raises Not_found (fun () -> SplayMap.find m (n/2));
+  ()
+
+let test_big () =
+  let n = 100_000 in
+  let m = SplayMap.of_seq (SplayMap.empty ())
+    (Sequence.zip (Sequence.zip_i (Sequence.int_range ~start:0 ~stop:n))) in
+  for i = 0 to n do
+    OUnit.assert_equal ~printer:string_of_int i (SplayMap.find m i);
+  done;
+  OUnit.assert_equal ~printer:string_of_int (n+1) (SplayMap.size m);
+  ()
+
+let suite =
+  "test_splayMap" >:::
+    [ "test1" >:: test1;
+      "test_remove" >:: test_remove;
+      "test_big" >:: test_big;
+    ]