diff --git a/splayMap.ml b/splayMap.ml index d0afedbe..457baf65 100644 --- a/splayMap.ml +++ b/splayMap.ml @@ -29,6 +29,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. http://www.cs.cornell.edu/Courses/cs3110/2009fa/recitations/rec-splay.html *) +(** {2 Polymorphic Maps} *) type ('a, 'b) t = { cmp : 'a -> 'a -> int; @@ -196,3 +197,216 @@ let to_seq t = let of_seq t seq = Sequence.fold (fun t (k, v) -> add t k v) t seq + +(** {2 Functorial interface} *) + +module type S = sig + type key + type 'a t + (** Tree with keys of type [key] and values of type 'a *) + + val empty : unit -> 'a t + (** Empty tree *) + + val is_empty : _ t -> bool + (** Is the tree empty? *) + + val find : 'a t -> key -> 'a + (** Find the element for this key, or raises Not_found *) + + val mem : _ t -> key -> bool + (** Is the key member of the tree? *) + + val add : 'a t -> key -> 'a -> 'a t + (** Add the binding to the tree *) + + val singleton : key -> 'a -> 'a t + (** Singleton map *) + + val remove : 'a t -> key -> 'a t + (** Remove the binding for this key *) + + val iter : 'a t -> (key -> 'a -> unit) -> unit + (** Iterate on bindings *) + + val fold : 'a t -> 'c -> ('c -> key -> 'a -> 'c) -> 'c + (** Fold on bindings *) + + val size : _ t -> int + (** Number of bindings (linear) *) + + val choose : 'a t -> (key * 'a) + (** Some binding, or raises Not_found *) + + val to_seq : 'a t -> (key * 'a) Sequence.t + + val of_seq : 'a t -> (key * 'a) Sequence.t -> 'a t +end + +module type ORDERED = sig + type t + val compare : t -> t -> int +end + +module Make(X : ORDERED) = struct + + type key = X.t + type 'a t = { + mutable tree : 'a tree; (* for lookups *) + } (** Tree with keys of type key, and values of type 'a *) + and 'a tree = + | Empty + | Node of (key * 'a * 'a tree * 'a tree) + + let empty () = + { 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 (k, v, l, r) key = + let c = X.compare 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 = X.compare 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 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 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 = X.compare 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 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 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 (k, v, l, r) key in + t.tree <- Node (k, v, l, r); (* save balanced tree *) + if X.compare 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 (k, v, l, r) key in + t.tree <- Node (k, v, l, r); (* save balanced tree *) + if X.compare key k = 0 + then true + else false + + (** Recursive insertion of key->value in the tree *) + let rec insert tree key value = + match tree with + | Empty -> Node (key, value, Empty, Empty) + | Node (k, v, l, r) -> + let c = X.compare key k in + if c = 0 + then Node (key, value, l, r) (* replace *) + else if c < 0 + then Node (k, v, insert l key value, r) + else Node (k, v, l, insert 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 (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 tree key value + in + { tree; } + + let singleton key value = + add (empty ()) 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 (k, v, l, r) key in + t.tree <- Node (k, v, l, r); + if X.compare key k = 0 + then (* remove the node, by merging the subnodes *) + let tree = left_merge l r in + { tree; } + else (* not present, same tree *) + t + + let iter t f = + let rec iter t = match t with + | Empty -> () + | Node (k, v, l, r) -> + iter l; + f k v; + 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 = fold acc l in + let acc = f acc k v 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 +end diff --git a/splayMap.mli b/splayMap.mli index f504da79..90f5d181 100644 --- a/splayMap.mli +++ b/splayMap.mli @@ -28,6 +28,8 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. (* TODO: map-wide operations: merge, compare, equal, for_all, exists, batch (sorted) add, partition, split, max_elt, min_elt, map... *) +(** {2 Polymorphic Maps} *) + type ('a, 'b) t (** Tree with keys of type 'a, and values of type 'b *) @@ -70,3 +72,55 @@ val choose : ('a, 'b) t -> ('a * 'b) val to_seq : ('a, 'b) t -> ('a * 'b) Sequence.t val of_seq : ('a, 'b) t -> ('a * 'b) Sequence.t -> ('a, 'b) t + +(** {2 Functorial interface} *) + +module type S = sig + type key + type 'a t + (** Tree with keys of type [key] and values of type 'a *) + + val empty : unit -> 'a t + (** Empty tree *) + + val is_empty : _ t -> bool + (** Is the tree empty? *) + + val find : 'a t -> key -> 'a + (** Find the element for this key, or raises Not_found *) + + val mem : _ t -> key -> bool + (** Is the key member of the tree? *) + + val add : 'a t -> key -> 'a -> 'a t + (** Add the binding to the tree *) + + val singleton : key -> 'a -> 'a t + (** Singleton map *) + + val remove : 'a t -> key -> 'a t + (** Remove the binding for this key *) + + val iter : 'a t -> (key -> 'a -> unit) -> unit + (** Iterate on bindings *) + + val fold : 'a t -> 'c -> ('c -> key -> 'a -> 'c) -> 'c + (** Fold on bindings *) + + val size : _ t -> int + (** Number of bindings (linear) *) + + val choose : 'a t -> (key * 'a) + (** Some binding, or raises Not_found *) + + val to_seq : 'a t -> (key * 'a) Sequence.t + + val of_seq : 'a t -> (key * 'a) Sequence.t -> 'a t +end + +module type ORDERED = sig + type t + val compare : t -> t -> int +end + +module Make(X : ORDERED) : S with type key = X.t