mirror of
https://github.com/c-cube/ocaml-containers.git
synced 2025-12-06 03:05:28 -05:00
leaner Heap
This commit is contained in:
Simon Cruanes
parent
35aa6cb03a
commit
02481a9352
1 changed files with 89 additions and 125 deletions
214
heap.ml
214
heap.ml
|
|
@ -25,140 +25,104 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|||
|
||||
(** {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)
|
||||
(** A splay tree with the given comparison function *)
|
||||
and 'a tree =
|
||||
| Empty
|
||||
| Node of ('a tree * 'a * 'a tree)
|
||||
(** A splay tree containing values of type 'a *)
|
||||
let empty ~cmp = {
|
||||
tree = Empty;
|
||||
cmp;
|
||||
}
|
||||
|
||||
let empty ~cmp =
|
||||
(Empty, cmp)
|
||||
let is_empty h =
|
||||
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
|
||||
| Empty -> true
|
||||
| Node _ -> false
|
||||
| Empty -> raise Not_found
|
||||
| Node (Empty, x, _) -> x
|
||||
| Node (l, _, _) -> min l
|
||||
in min h.tree
|
||||
|
||||
(** Partition the tree into (elements <= pivot, elements > pivot) *)
|
||||
let rec partition ~cmp pivot tree =
|
||||
(** Get minimum value and remove it from the 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
|
||||
| Empty -> 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
|
||||
iter a; f x; iter b
|
||||
in iter h.tree
|
||||
|
||||
(** Insert the element in the tree *)
|
||||
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 size h =
|
||||
let r = ref 0 in
|
||||
iter heap (fun _ -> incr r);
|
||||
iter h (fun _ -> incr r);
|
||||
!r
|
||||
|
||||
let to_seq heap =
|
||||
fun k -> iter heap k
|
||||
let to_seq h =
|
||||
fun k -> iter h k
|
||||
|
||||
let of_seq heap seq =
|
||||
seq (fun elt -> insert heap elt)
|
||||
let of_seq h seq =
|
||||
seq (fun elt -> insert h elt)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue