Skip to content

Commit

Permalink
Add remove_min_exn function.
Browse files Browse the repository at this point in the history
  • Loading branch information
lyrm committed Jan 8, 2025
1 parent 3a0f0fb commit 78c36ab
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 13 deletions.
28 changes: 16 additions & 12 deletions src/priority_queue.ml
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ type 'k compare = 'k -> 'k -> int

type ('k, 'v) t = { compare : 'k compare; root : ('k, 'v) links; size : Size.t }

exception Empty
(* *)

(** [get_random_height max_height] gives a random value [n] in the range from
Expand Down Expand Up @@ -293,6 +294,8 @@ let length t = Size.get t.size

(* *)

type ('a, _) poly = Option : ('a, 'a option) poly | Value : ('a, 'a) poly

let rec find_min t : (_, _, [< `Node | `Null ]) node =
let root = t.root in
let root_at_level0 = Array.unsafe_get root 0 in
Expand Down Expand Up @@ -336,15 +339,16 @@ let rec try_remove t key next level link = function
try_remove t key next level link (Atomic.get link)
else try_remove t key next level link (Atomic.get link)

let remove_min_opt t =
let rec loop t =
match find_min t with
| Null -> None
| Node { next; key; value; _ } ->
let level = Array.length next - 1 in
let link = Array.unsafe_get next level in
if try_remove t key next level link (Atomic.get link) then
Some (key, value)
else loop t
in
loop t
let rec remove_min_as : type p v r. (p, v) t -> (p * v, r) poly -> r =
fun t poly ->
match find_min t with
| Null -> ( match poly with Value -> raise Empty | Option -> None)
| Node { next; key; value; _ } ->
let level = Array.length next - 1 in
let link = Array.unsafe_get next level in
if try_remove t key next level link (Atomic.get link) then
match poly with Value -> (key, value) | Option -> Some (key, value)
else remove_min_as t poly

let remove_min_opt t = remove_min_as t Option
let remove_min_exn t = remove_min_as t Value
11 changes: 10 additions & 1 deletion src/priority_queue.mli
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,18 @@ val add : ('p, 'v) t -> 'p -> 'v -> unit
val remove_min_opt : ('p, 'v) t -> ('p * 'v) option
(** [remove_min_opt s] removes and returns [Some] of the element with the
smallest priority from the priority queue [pq]. If the priority queue [pq]
was empty [None] is returned. The elements with the same priority are
was empty, [None] is returned. The elements with the same priority are
removed in fifo order. *)

exception Empty

val remove_min_exn : ('p, 'v) t -> 'p * 'v
(** [remove_min_exn s] removes and returns the element with the smallest
priority from the priority queue [pq]. The elements with the same priority
are removed in fifo order.
@raise Empty if the priority queue [pq] is empty.*)

val length : ('k, 'v) t -> int
(** [length pq] computes and returns the number of elements in the priority
queue [pq].
Expand Down

0 comments on commit 78c36ab

Please sign in to comment.