mirror of
https://github.com/c-cube/ocaml-containers.git
synced 2025-12-06 03:05:28 -05:00
functional SplayTree as Maps
This commit is contained in:
Simon Cruanes
parent
d0bb88e365
commit
6c1f7fb950
4 changed files with 309 additions and 0 deletions
198
splayMap.ml
Normal file
198
splayMap.ml
Normal file
|
|
@ -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
|
||||
69
splayMap.mli
Normal file
69
splayMap.mli
Normal file
|
|
@ -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
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
41
tests/test_splayMap.ml
Normal file
41
tests/test_splayMap.ml
Normal file
|
|
@ -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;
|
||||
]
|
||||
Loading…
Add table
Reference in a new issue