leaner Heap

This commit is contained in:
Simon Cruanes 2013-03-20 15:54:45 +01:00
parent 35aa6cb03a
commit 02481a9352

110
heap.ml
View file

@ -25,25 +25,27 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
(** {1 Imperative priority queue} *) (** {1 Imperative priority queue} *)
module Tree = struct type 'a t = {
mutable tree : 'a tree;
type 'a t = 'a tree * ('a -> 'a -> int) cmp : 'a -> 'a -> int;
(** A splay tree with the given comparison function *) } (** A splay tree heap with the given comparison function *)
and 'a tree = and 'a tree =
| Empty | Empty
| Node of ('a tree * 'a * 'a tree) | Node of ('a tree * 'a * 'a tree)
(** A splay tree containing values of type 'a *) (** A splay tree containing values of type 'a *)
let empty ~cmp = let empty ~cmp = {
(Empty, cmp) tree = Empty;
cmp;
}
let is_empty (tree, _) = let is_empty h =
match tree with match h.tree with
| Empty -> true | Empty -> true
| Node _ -> false | Node _ -> false
(** Partition the tree into (elements <= pivot, elements > pivot) *) (** Partition the tree into (elements <= pivot, elements > pivot) *)
let rec partition ~cmp pivot tree = let rec partition ~cmp pivot tree =
match tree with match tree with
| Empty -> Empty, Empty | Empty -> Empty, Empty
| Node (a, x, b) -> | Node (a, x, b) ->
@ -72,29 +74,23 @@ module Tree = struct
small, Node (big, y, Node (a2, x, b)) small, Node (big, y, Node (a2, x, b))
end end
(** Insert the element in the tree *) (** Insert the element in the tree *)
let insert (tree, cmp) x = let insert h x =
let small, big = partition ~cmp x tree in let small, big = partition ~cmp:h.cmp x h.tree in
let tree' = Node (small, x, big) in let tree' = Node (small, x, big) in
tree', cmp h.tree <- tree'
(** Returns the top value, or raise Not_found is empty *) (** Access minimum value *)
let top (tree, _) = let min h =
match tree with
| Empty -> raise Not_found
| Node (_, x, _) -> x
(** Access minimum value *)
let min (tree, _) =
let rec min tree = let rec min tree =
match tree with match tree with
| Empty -> raise Not_found | Empty -> raise Not_found
| Node (Empty, x, _) -> x | Node (Empty, x, _) -> x
| Node (l, _, _) -> min l | Node (l, _, _) -> min l
in min tree in min h.tree
(** Get minimum value and remove it from the tree *) (** Get minimum value and remove it from the tree *)
let delete_min (tree, cmp) = let pop h =
let rec delete_min tree = match tree with let rec delete_min tree = match tree with
| Empty -> raise Not_found | Empty -> raise Not_found
| Node (Empty, x, b) -> x, b | Node (Empty, x, b) -> x, b
@ -104,61 +100,29 @@ module Tree = struct
let m, a' = delete_min a in let m, a' = delete_min a in
m, Node (a', x, Node (b, y, c)) m, Node (a', x, Node (b, y, c))
in in
let m, tree' = delete_min tree in let m, tree' = delete_min h.tree in
m, (tree', cmp) h.tree <- tree';
m
(** Iterate on elements *) let junk h =
let iter (tree, _) f = ignore (pop h)
(** Iterate on elements *)
let iter h f =
let rec iter tree = let rec iter tree =
match tree with match tree with
| Empty -> () | Empty -> ()
| Node (a, x, b) -> | Node (a, x, b) ->
iter a; f x; iter b iter a; f x; iter b
in iter tree in iter h.tree
end
type 'a t = 'a Tree.t ref let size h =
(** The heap is a reference to a splay tree *)
(** Create an empty heap *)
let empty ~cmp =
ref (Tree.empty ~cmp)
(** Insert a value in the heap *)
let insert heap x =
heap := Tree.insert !heap x
(** Check whether the heap is empty *)
let is_empty heap =
Tree.is_empty !heap
(** Access the minimal value of the heap, or raises Empty *)
let min (heap : 'a t) : 'a =
let elt = Tree.min !heap in
elt
(** Discard the minimal element *)
let junk heap =
let _, tree' = Tree.delete_min !heap in
heap := tree'
(** Remove and return the mininal value (or raise Invalid_argument) *)
let pop heap =
let elt, tree' = Tree.delete_min !heap in
heap := tree';
elt
(** Iterate on the elements, in an unspecified order *)
let iter heap k =
Tree.iter !heap (fun elt -> k elt)
let size heap =
let r = ref 0 in let r = ref 0 in
iter heap (fun _ -> incr r); iter h (fun _ -> incr r);
!r !r
let to_seq heap = let to_seq h =
fun k -> iter heap k fun k -> iter h k
let of_seq heap seq = let of_seq h seq =
seq (fun elt -> insert heap elt) seq (fun elt -> insert h elt)