diff options
Diffstat (limited to 'src/tree/splay.ml')
-rw-r--r-- | src/tree/splay.ml | 194 |
1 files changed, 194 insertions, 0 deletions
diff --git a/src/tree/splay.ml b/src/tree/splay.ml new file mode 100644 index 0000000..662fc6c --- /dev/null +++ b/src/tree/splay.ml @@ -0,0 +1,194 @@ +module type KEY = sig + + type 'a t + + (** Parametrized comparator *) + val comp: 'a t -> 'b t -> ('a, 'b) Tools.cmp + + val repr: Format.formatter -> 'a t -> unit + +end + +module Make (El : KEY) = struct + + type container = C : ('a El.t * 'a) -> container [@@unboxed] + + type leaf (** Fantom type for typing the tree *) + type node (** Fantom type for typing the tree *) + + type 'a branch = + | Leaf : leaf branch + | Node : _ branch * ('a El.t * 'a) * _ branch -> node branch + + type t = T : 'a branch ref -> t [@@unboxed] + + let empty = T (ref Leaf) + + let isEmpty (T tree) = match !tree with + | Leaf -> true + | _ -> false + + let rec splay : type a. a El.t -> node branch -> node branch = fun x t -> begin + let Node (l, y, r) = t in + begin match El.comp x (fst y) with + | Tools.Eq -> t + | Tools.Lt -> + begin match l with + | Leaf -> t + | Node (ll, z, rr) -> + begin match El.comp x (fst z) with + | Tools.Eq -> Node (ll, z, Node (rr, y, r)) + | Tools.Lt -> + begin match ll with + | Leaf -> Node (ll, z, Node (rr, y, r)) + | Node _ as ll -> + let Node (newL, newV, newR) = splay x ll + in Node (newL, newV, Node (newR, z, Node (rr, y, r))) + end + | Tools.Gt -> + begin match rr with + | Leaf -> Node (ll, z, Node (rr, y, r)) + | Node _ as rr -> + let Node (newL, newV, newR) = splay x rr + in Node (Node (ll, z, newL), newV, Node (newR, y, r)) + end + end + end + | Tools.Gt -> + begin match r with + | Leaf -> t + | Node (ll, z, rr) -> + begin match El.comp x (fst z) with + | Tools.Eq -> Node (Node (l, y, ll), z, rr) + | Tools.Lt -> + begin match ll with + | Leaf -> Node (Node (l, y, ll), z, rr) + | Node _ as ll -> + let Node (newL, newV, newR) = splay x ll + in Node (Node (l, y, newL), newV, Node (newR, z, rr)) + end + | Tools.Gt -> + begin match rr with + | Leaf -> Node (Node (l, y, ll), z, rr) + | Node _ as rr -> + let Node (newL, newV, newR) = splay x rr + in Node (Node (Node(l, y, ll), z, newL), newV, newR) + end + end + end + end + end + + let member: type a. a El.t -> t -> bool = fun x (T t) -> match !t with + | Leaf -> false + | Node _ as root -> + let root' = splay x root in + t := root'; + let Node (_, c', _) = root' in + begin match El.comp (fst c') x with + | Tools.Eq -> true + | _ -> false + end + + let find: type a. a El.t -> t -> a = fun x (T t) -> match !t with + | Leaf -> raise Not_found + | Node _ as root -> + let root' = splay x root in + t := root'; + let Node (_, c', _) = root' in + begin match El.comp (fst c') x with + | Tools.Eq -> snd c' + | _ -> raise Not_found + end + + let add: type a. a El.t -> a -> t -> t = fun key value (T t) -> match !t with + | Leaf -> T (ref (Node (Leaf, (key, value), Leaf))) + | Node _ as root -> + let root' = splay key root in + let Node (l, y, r) = root' in + begin match El.comp key (fst y) with + | Tools.Eq -> T (ref (Node(l, (key, value), r))) + | Tools.Lt -> T (ref (Node (l, (key, value), Node (Leaf, y, r)))) + | Tools.Gt -> T (ref (Node (Node (l, y, Leaf), (key, value), r))) + end + + let rec _subtree_maximum:type a. a branch -> a branch = fun t -> begin match t with + | Leaf -> Leaf + | Node (_, _, (Node (_, _, _) as x)) -> _subtree_maximum x + | Node (_, (key, value), Leaf) -> splay key t + end + + let rec _subtree_minimum: type a. a branch -> a branch = fun t -> begin match t with + | Leaf -> Leaf + | Node ((Node (_, _, _) as x), _, _) -> _subtree_minimum x + | Node (Leaf, (key, value), _) -> splay key t + end + + let remove: type a. a El.t -> t -> t = fun key (T t) -> begin match !t with + | Leaf -> empty + | Node _ as root -> + let root' = splay key root in + let Node (l, c', r) = root' in + begin match El.comp (fst c') key with + | Tools.Eq -> begin match _subtree_maximum l with + | Node(l, c, Leaf) -> T (ref (Node(l, c, r))) + | Node(l, c, _) -> raise Not_found + | Leaf -> begin match _subtree_minimum r with + | Leaf -> empty + | Node(Leaf, c, r) -> T (ref (Node(l, c, r))) + | Node(_, c, r) -> raise Not_found + end + end + (* The key is not present, return the splayed tree *) + | _ -> T (ref root') + end + end + + (** Existencial type for the branches *) + type exBranch = Branch : _ branch -> exBranch [@@unboxed] + + let fold f init (T t) = begin + let rec _fold : type b. (container * exBranch) list -> 'a -> b branch -> 'a = begin + fun acc v -> function + (* We have a node : we accumulate the right part, and process the left branch *) + | Node (left, (key, value), right) -> + let c = C (key, value) in + (_fold [@tailcall]) ((c, Branch right)::acc) v left + (* We have nothing left, we process the values delayed *) + | Leaf -> begin match acc with + | [] -> v + | (c, (Branch right))::tl -> (_fold [@tailcall]) tl (f v c) right + end + end in + _fold [] init !t + end + + let repr formatter (T t) = begin + + let repr_edge from formatter dest = begin + Format.fprintf formatter "\"%a\" -> \"%a\"\n" + El.repr from + El.repr dest + end in + + let rec repr': type a b. a El.t -> Format.formatter -> b branch -> unit = fun parent formatter -> function + | Leaf -> () + | Node (l, c, r) -> + let key = fst c in + Format.fprintf formatter "%a%a%a" + (repr_edge parent) key + (repr' key) l + (repr' key) r in + + begin match !t with + | Leaf -> Format.fprintf formatter "digraph G {}" + | Node (l, c, r) -> + let key = fst c in + Format.fprintf formatter "digraph G {\n%a%a}" + (repr' key) l + (repr' key) r + end + + end + +end |