1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
open! Core
module Array = Base.Array
module List = Base.List
module Option = Base.Option
module Sequence = Base.Sequence
module Node = struct
type 'a t =
{ value : 'a
; children : 'a t list
}
end
open Node
type 'a t =
{ compare : 'a -> 'a -> int
; length : int
; heap : 'a Node.t option
}
let create ~compare = { compare; length = 0; heap = None }
let merge
~compare
({ value = e1; children = nl1 } as n1)
({ value = e2; children = nl2 } as n2)
=
if compare e1 e2 < 0
then { value = e1; children = n2 :: nl1 }
else { value = e2; children = n1 :: nl2 }
;;
let merge_pairs ~compare t =
let rec loop acc t =
match t with
| [] -> acc
| [ head ] -> head :: acc
| head :: next1 :: next2 -> loop (merge ~compare head next1 :: acc) next2
in
match loop [] t with
| [] -> None
| [ h ] -> Some h
| x :: xs -> Some (List.fold xs ~init:x ~f:(merge ~compare))
;;
let add { compare; length; heap } e =
let new_node = { value = e; children = [] } in
let heap =
match heap with
| None -> new_node
| Some heap -> merge ~compare new_node heap
in
{ compare; length = length + 1; heap = Some heap }
;;
let top_exn t =
match t.heap with
| None -> failwith "Fheap.top_exn called on an empty heap"
| Some { value; _ } -> value
;;
let top t =
try Some (top_exn t) with
| _ -> None
;;
let pop_exn { compare; length; heap } =
match heap with
| None -> failwith "Heap.pop_exn called on an empty heap"
| Some { value; children } ->
let new_heap = merge_pairs ~compare children in
let t' = { compare; length = length - 1; heap = new_heap } in
value, t'
;;
let pop t =
try Some (pop_exn t) with
| _ -> None
;;
let pop_min = pop
let pop_min_exn = pop_exn
let remove_top t =
try
let _, t' = pop_exn t in
Some t'
with
| _ -> None
;;
let pop_if t f =
match top t with
| None -> None
| Some v -> if f v then pop t else None
;;
let fold t ~init ~f =
let rec loop acc to_visit =
match to_visit with
| [] -> acc
| { value; children } :: rest ->
let acc = f acc value in
let to_visit = List.unordered_append children rest in
loop acc to_visit
in
match t.heap with
| None -> init
| Some node -> loop init [ node ] [@nontail]
;;
let length t = t.length
module C = Container.Make (struct
type nonrec 'a t = 'a t
let fold = fold
let iter = `Define_using_fold
let length = `Custom length
end)
let is_empty t = Option.is_none t.heap
let iter = C.iter
let mem = C.mem
let min_elt = C.min_elt
let max_elt = C.max_elt
let find = C.find
let find_map = C.find_map
let for_all = C.for_all
let exists = C.exists
let sum = C.sum
let count = C.count
let to_list = C.to_list
let fold_result = C.fold_result
let fold_until = C.fold_until
let to_array = C.to_array
type ('a, 'at, 'accum) folder = 'at -> init:'accum -> f:('accum -> 'a -> 'accum) -> 'accum
let of_fold c ~compare (fold : _ folder) =
let h = create ~compare in
fold c ~init:h ~f:add
;;
let of_list l ~compare = of_fold l ~compare List.fold
let of_array arr ~compare = of_fold arr ~compare Array.fold
let sexp_of_t sexp_of_a t = List.sexp_of_t sexp_of_a (to_list t)
let to_sequence t = Sequence.unfold ~init:t ~f:pop