simon/ocaml-containers
1
0
Fork 0
mirror of https://github.com/c-cube/ocaml-containers.git synced 2025-12-06 11:15:31 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

leaner Heap

Browse source
This commit is contained in:
Simon Cruanes 2013-03-20 15:54:45 +01:00
parent 35aa6cb03a
commit 02481a9352
1 changed files with 89 additions and 125 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

214
heap.ml
View file

@ -25,140 +25,104 @@ 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;
cmp : 'a -> 'a -> int;
} (** A splay tree heap with the given comparison function *)
and 'a tree =
| Empty
| Node of ('a tree * 'a * 'a tree)
(** A splay tree containing values of type 'a *)
type 'a t = 'a tree * ('a -> 'a -> int) let empty ~cmp = {
(** A splay tree with the given comparison function *) tree = Empty;
and 'a tree = cmp;
| Empty }
| Node of ('a tree * 'a * 'a tree)
(** A splay tree containing values of type 'a *)
let empty ~cmp = let is_empty h =
(Empty, cmp) match h.tree with
| Empty -> true
| Node _ -> false
let is_empty (tree, _) = (** Partition the tree into (elements <= pivot, elements > pivot) *)
let rec partition ~cmp pivot tree =
match tree with
| Empty -> Empty, Empty
| Node (a, x, b) ->
if cmp x pivot <= 0
then begin
match b with
| Empty -> (tree, Empty)
| Node (b1, y, b2) ->
if cmp y pivot <= 0
then
let small, big = partition ~cmp pivot b2 in
Node (Node (a, x, b1), y, small), big
else
let small, big = partition ~cmp pivot b1 in
Node (a, x, small), Node (big, y, b2)
end else begin
match a with
| Empty -> (Empty, tree)
| Node (a1, y, a2) ->
if cmp y pivot <= 0
then
let small, big = partition ~cmp pivot a2 in
Node (a1, y, small), Node (big, x, b)
else
let small, big = partition ~cmp pivot a1 in
small, Node (big, y, Node (a2, x, b))
end
(** Insert the element in the tree *)
let insert h x =
let small, big = partition ~cmp:h.cmp x h.tree in
let tree' = Node (small, x, big) in
h.tree <- tree'
(** Access minimum value *)
let min h =
let rec min tree =
match tree with match tree with
| Empty -> true | Empty -> raise Not_found
| Node _ -> false | Node (Empty, x, _) -> x
| Node (l, _, _) -> min l
in min h.tree
(** Partition the tree into (elements <= pivot, elements > pivot) *) (** Get minimum value and remove it from the tree *)
let rec partition ~cmp pivot tree = let pop h =
let rec delete_min tree = match tree with
| Empty -> raise Not_found
| Node (Empty, x, b) -> x, b
| Node (Node (Empty, x, b), y, c) ->
x, Node (b, y, c) (* rebalance *)
| Node (Node (a, x, b), y, c) ->
let m, a' = delete_min a in
m, Node (a', x, Node (b, y, c))
in
let m, tree' = delete_min h.tree in
h.tree <- tree';
m
let junk h =
ignore (pop h)
(** Iterate on elements *)
let iter h f =
let rec iter tree =
match tree with match tree with
| Empty -> Empty, Empty | Empty -> ()
| Node (a, x, b) -> | Node (a, x, b) ->
if cmp x pivot <= 0 iter a; f x; iter b
then begin in iter h.tree
match b with
| Empty -> (tree, Empty)
| Node (b1, y, b2) ->
if cmp y pivot <= 0
then
let small, big = partition ~cmp pivot b2 in
Node (Node (a, x, b1), y, small), big
else
let small, big = partition ~cmp pivot b1 in
Node (a, x, small), Node (big, y, b2)
end else begin
match a with
| Empty -> (Empty, tree)
| Node (a1, y, a2) ->
if cmp y pivot <= 0
then
let small, big = partition ~cmp pivot a2 in
Node (a1, y, small), Node (big, x, b)
else
let small, big = partition ~cmp pivot a1 in
small, Node (big, y, Node (a2, x, b))
end
(** Insert the element in the tree *) let size h =
let insert (tree, cmp) x =
let small, big = partition ~cmp x tree in
let tree' = Node (small, x, big) in
tree', cmp
(** Returns the top value, or raise Not_found is empty *)
let top (tree, _) =
match tree with
| Empty -> raise Not_found
| Node (_, x, _) -> x
(** Access minimum value *)
let min (tree, _) =
let rec min tree =
match tree with
| Empty -> raise Not_found
| Node (Empty, x, _) -> x
| Node (l, _, _) -> min l
in min tree
(** Get minimum value and remove it from the tree *)
let delete_min (tree, cmp) =
let rec delete_min tree = match tree with
| Empty -> raise Not_found
| Node (Empty, x, b) -> x, b
| Node (Node (Empty, x, b), y, c) ->
x, Node (b, y, c) (* rebalance *)
| Node (Node (a, x, b), y, c) ->
let m, a' = delete_min a in
m, Node (a', x, Node (b, y, c))
in
let m, tree' = delete_min tree in
m, (tree', cmp)
(** Iterate on elements *)
let iter (tree, _) f =
let rec iter tree =
match tree with
| Empty -> ()
| Node (a, x, b) ->
iter a; f x; iter b
in iter tree
end
type 'a t = 'a Tree.t ref
(** 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)