Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Newer
Older
100644 391 lines (327 sloc) 11.599 kb
fccc685 Initial open-source release
MLstate authored
1 (*
2 Copyright © 2011 MLstate
3
4 This file is part of OPA.
5
6 OPA is free software: you can redistribute it and/or modify it under the
7 terms of the GNU Affero General Public License, version 3, as published by
8 the Free Software Foundation.
9
10 OPA is distributed in the hope that it will be useful, but WITHOUT ANY
11 WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
12 FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for
13 more details.
14
15 You should have received a copy of the GNU Affero General Public License
16 along with OPA. If not, see <http://www.gnu.org/licenses/>.
17 *)
18 module Make (Ord: OrderedTypeSig.S) : (BaseSetSig.S with type elt = Ord.t) =
19 struct
20 type elt = Ord.t
21 type t =
22 | Empty
23 | Node of t * elt * t * int (* int (* size *) *)
24
25 (* Sets are represented by balanced binary trees (the heights of the
26 children differ by at most 2 *)
27
28 let height = function
29 Empty -> 0
30 | Node(_, _, _, h) -> h
31
32 (* Creates a new node with left son l, value v and right son r.
33 We must have all elements of l < v < all elements of r.
34 l and r must be balanced and | height l - height r | <= 2.
35 Inline expansion of height for better speed. *)
36
37 let create l v r =
38 let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
39 let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
40 Node(l, v, r, (if hl >= hr then hl + 1 else hr + 1))
41
42 (* Same as create, but performs one step of rebalancing if necessary.
43 Assumes l and r balanced and | height l - height r | <= 3.
44 Inline expansion of create for better speed in the most frequent case
45 where no rebalancing is required. *)
46
47 let bal l v r =
48 let hl = match l with Empty -> 0 | Node(_,_,_,h) -> h in
49 let hr = match r with Empty -> 0 | Node(_,_,_,h) -> h in
50 if hl > hr + 2 then begin
51 match l with
52 Empty -> invalid_arg "Set.bal"
53 | Node(ll, lv, lr, _) ->
54 if height ll >= height lr then
55 create ll lv (create lr v r)
56 else begin
57 match lr with
58 Empty -> invalid_arg "Set.bal"
59 | Node(lrl, lrv, lrr, _)->
60 create (create ll lv lrl) lrv (create lrr v r)
61 end
62 end else if hr > hl + 2 then begin
63 match r with
64 Empty -> invalid_arg "Set.bal"
65 | Node(rl, rv, rr, _) ->
66 if height rr >= height rl then
67 create (create l v rl) rv rr
68 else begin
69 match rl with
70 Empty -> invalid_arg "Set.bal"
71 | Node(rll, rlv, rlr, _) ->
72 create (create l v rll) rlv (create rlr rv rr)
73 end
74 end else
75 Node(l, v, r, (if hl >= hr then hl + 1 else hr + 1))
76
77 (* Insertion of one element *)
78
79 let rec add x = function
80 Empty -> Node(Empty, x, Empty, 1)
81 | Node(l, v, r, _) as t ->
82 let c = Ord.compare x v in
83 if c = 0 then t else
84 if c < 0 then bal (add x l) v r else bal l v (add x r)
85
86 (* Same as create and bal, but no assumptions are made on the
87 relative heights of l and r. *)
88
89 let rec join l v r =
90 match (l, r) with
91 (Empty, _) -> add v r
92 | (_, Empty) -> add v l
93 | (Node(ll, lv, lr, lh), Node(rl, rv, rr, rh)) ->
94 if lh > rh + 2 then bal ll lv (join lr v r) else
95 if rh > lh + 2 then bal (join l v rl) rv rr else
96 create l v r
97
98 (* Smallest and greatest element of a set *)
99
100 let rec min_elt = function
101 Empty -> raise Not_found
102 | Node(Empty, v, _r, _) -> v
103 | Node(l, _v, _r, _) -> min_elt l
104
105 let rec max_elt = function
106 Empty -> raise Not_found
107 | Node(_l, v, Empty, _) -> v
108 | Node(_l, _v, r, _) -> max_elt r
109
110 (* Remove the smallest element of the given set *)
111
112 let rec remove_min_elt = function
113 Empty -> invalid_arg "Set.remove_min_elt"
114 | Node(Empty, _v, r, _) -> r
115 | Node(l, v, r, _) -> bal (remove_min_elt l) v r
116
117 (* Merge two trees l and r into one.
118 All elements of l must precede the elements of r.
119 Assume | height l - height r | <= 2. *)
120
121 let merge t1 t2 =
122 match (t1, t2) with
123 (Empty, t) -> t
124 | (t, Empty) -> t
125 | (_, _) -> bal t1 (min_elt t2) (remove_min_elt t2)
126
127 (* Merge two trees l and r into one.
128 All elements of l must precede the elements of r.
129 No assumption on the heights of l and r. *)
130
131 let concat t1 t2 =
132 match (t1, t2) with
133 (Empty, t) -> t
134 | (t, Empty) -> t
135 | (_, _) -> join t1 (min_elt t2) (remove_min_elt t2)
136
137 (* Splitting. split x s returns a triple (l, present, r) where
138 - l is the set of elements of s that are < x
139 - r is the set of elements of s that are > x
140 - present is false if s contains no element equal to x,
141 or true if s contains an element equal to x. *)
142
143 let rec split x = function
144 Empty ->
145 (Empty, false, Empty)
146 | Node(l, v, r, _) ->
147 let c = Ord.compare x v in
148 if c = 0 then (l, true, r)
149 else if c < 0 then
150 let (ll, pres, rl) = split x l in (ll, pres, join rl v r)
151 else
152 let (lr, pres, rr) = split x r in (join l v lr, pres, rr)
153
154 (* Implementation of the set operations *)
155
156 let empty = Empty
157
158 let is_empty = function Empty -> true | _ -> false
159
160 let rec mem x = function
161 Empty -> false
162 | Node(l, v, r, _) ->
163 let c = Ord.compare x v in
164 c = 0 || mem x (if c < 0 then l else r)
165
166 let singleton x = Node(Empty, x, Empty, 1)
167
168 let rec remove x = function
169 Empty -> Empty
170 | Node(l, v, r, _) ->
171 let c = Ord.compare x v in
172 if c = 0 then merge l r else
173 if c < 0 then bal (remove x l) v r else bal l v (remove x r)
174
175 let rec size = function
176 | Empty -> 0
177 | Node (l, _, r, _) -> 1 + size l + size r
178
179 let draw t =
180 let rec aux = function
181 | Empty -> raise Not_found
182 | Node (l, v, r, _) ->
183 let sl = size l
184 and sr = size r in
185 match Random.int (1 + sl + sr) with
186 | 0 -> v, remove v t
187 | i when i <= sl -> aux l
188 | _ -> aux r
189 in
190 aux t
191
192 let rec union s1 s2 =
193 match (s1, s2) with
194 (Empty, t2) -> t2
195 | (t1, Empty) -> t1
196 | (Node(l1, v1, r1, h1), Node(l2, v2, r2, h2)) ->
197 if h1 >= h2 then
198 if h2 = 1 then add v2 s1 else begin
199 let (l2, _, r2) = split v1 s2 in
200 join (union l1 l2) v1 (union r1 r2)
201 end
202 else
203 if h1 = 1 then add v1 s2 else begin
204 let (l1, _, r1) = split v2 s1 in
205 join (union l1 l2) v2 (union r1 r2)
206 end
207
208 let rec inter s1 s2 =
209 match (s1, s2) with
210 (Empty, _t2) -> Empty
211 | (_t1, Empty) -> Empty
212 | (Node(l1, v1, r1, _), t2) ->
213 match split v1 t2 with
214 (l2, false, r2) ->
215 concat (inter l1 l2) (inter r1 r2)
216 | (l2, true, r2) ->
217 join (inter l1 l2) v1 (inter r1 r2)
218
219 let rec diff s1 s2 =
220 match (s1, s2) with
221 (Empty, _t2) -> Empty
222 | (t1, Empty) -> t1
223 | (Node(l1, v1, r1, _), t2) ->
224 match split v1 t2 with
225 (l2, false, r2) ->
226 join (diff l1 l2) v1 (diff r1 r2)
227 | (l2, true, r2) ->
228 concat (diff l1 l2) (diff r1 r2)
229
230 type enumeration = End | More of elt * t * enumeration
231
232 let rec cons_enum s e =
233 match s with
234 Empty -> e
235 | Node(l, v, r, _) -> cons_enum l (More(v, r, e))
236
237 let rec compare_aux e1 e2 =
238 match (e1, e2) with
239 (End, End) -> 0
240 | (End, _) -> -1
241 | (_, End) -> 1
242 | (More(v1, r1, e1), More(v2, r2, e2)) ->
243 let c = Ord.compare v1 v2 in
244 if c <> 0
245 then c
246 else compare_aux (cons_enum r1 e1) (cons_enum r2 e2)
247
248 let compare s1 s2 =
249 compare_aux (cons_enum s1 End) (cons_enum s2 End)
250
251 let equal s1 s2 =
252 compare s1 s2 = 0
253
254 let rec subset s1 s2 =
255 match (s1, s2) with
256 Empty, _ ->
257 true
258 | _, Empty ->
259 false
260 | Node (l1, v1, r1, _), (Node (l2, v2, r2, _) as t2) ->
261 let c = Ord.compare v1 v2 in
262 if c = 0 then
263 subset l1 l2 && subset r1 r2
264 else if c < 0 then
265 subset (Node (l1, v1, Empty, 0)) l2 && subset r1 t2
266 else
267 subset (Node (Empty, v1, r1, 0)) r2 && subset l1 t2
268
269 let rec iter f = function
270 Empty -> ()
271 | Node(l, v, r, _) -> iter f l; f v; iter f r
272
273 let rec fold f s accu =
274 match s with
275 Empty -> accu
276 | Node(l, v, r, _) -> fold f r (f v (fold f l accu))
277
278 let rec fold_rev f s accu =
279 match s with
280 | Empty -> accu
281 | Node(l, v, r, _) -> fold_rev f l (f v (fold_rev f r accu))
282
283 let map f s =
284 fold (fun x acc -> add (f x) acc) s Empty
285
286 let rec for_all p = function
287 Empty -> true
288 | Node(l, v, r, _) -> p v && for_all p l && for_all p r
289
290 let rec exists p = function
291 Empty -> false
292 | Node(l, v, r, _) -> p v || exists p l || exists p r
293
294 let filter p s =
295 let rec filt accu = function
296 | Empty -> accu
297 | Node(l, v, r, _) ->
298 filt (filt (if p v then add v accu else accu) l) r in
299 filt Empty s
300
301 let partition p s =
302 let rec part (t, f as accu) = function
303 | Empty -> accu
304 | Node(l, v, r, _) ->
305 part (part (if p v then (add v t, f) else (t, add v f)) l) r in
306 part (Empty, Empty) s
307
308 let rec cardinal = function
309 Empty -> 0
310 | Node(l, _v, r, _) -> cardinal l + 1 + cardinal r
311
312 let rec elements_aux accu = function
313 Empty -> accu
314 | Node(l, v, r, _) -> elements_aux (v :: elements_aux accu r) l
315
316 let elements s =
317 elements_aux [] s
318
319 let add_list l t = List.fold_left (fun acc v -> add v acc) t l
320 let from_list l = add_list l empty
321
6f7b0d3 [enhance] Set: choose become quicker
Arthur Milchior authored
322 let choose = function
323 | Empty -> raise Not_found
324 | Node(_, v, _, _) -> v
fccc685 Initial open-source release
MLstate authored
325
e74f6b5 [feature] libbase: choose_opt in Maps and Sets
Arthur Milchior authored
326 let rec choose_opt = function
327 | Empty -> None
328 | Node (_, v, _r, _) -> Some v
329
645f07a [feature] libbase: Taking an element to show the difference between …
Arthur Milchior authored
330 let example_diff s1 s2 =
331 let diff_ = diff s1 s2 in
332 match choose_opt diff_ with
333 | Some elt -> Some elt
334 | None ->
335 let diff = diff s2 s1 in
336 match choose_opt diff with
337 | Some elt -> Some elt
338 | None -> None
339
fccc685 Initial open-source release
MLstate authored
340 let complete_join big small =
341 fold add small big
342
343 let rec complete fun_prefixe k = function
344 | Empty -> Empty
345 | Node (l, key, r, _) ->
346 if (fun_prefixe k key) then
347 (** k est un prefixe de key *)
348 let set_1 = complete fun_prefixe k l in
349 let set_2 = complete fun_prefixe k r in
350 let joined_set =
351 if (height set_1) >= (height set_2)
352 then complete_join set_1 set_2
353 else complete_join set_2 set_1
354 in add key joined_set
355 else
356 if k < key then complete fun_prefixe k l
357 else complete fun_prefixe k r
358
359 (* cf doc *)
360 let pp sep ppe fmt t =
361 let fiter elt =
362 ppe fmt elt ;
363 Format.fprintf fmt sep
364 in
365 iter fiter t
366
367 let compare_elt = Ord.compare
368
369 let safe_union s1 s2 =
370 let u = union s1 s2 in
371 (* We ensure that the 2 sets to join were disjoint. This is the case if
372 the sum of their sizes equal the size of the set obtained after
373 union. *)
374 if not (size u = size s1 + size s2) then
375 raise (Invalid_argument "Base.Set.safe_union") ;
376 u
377
378 let from_sorted_array elts =
379 let rec aux left right =
380 if left > right
381 then Empty
382 else
383 let midle = (left + right) lsr 1 in
384 let left_tree = aux left (pred midle) in
385 let right_tree = aux (succ midle) right in
386 let elt = Array.unsafe_get elts midle in
387 create left_tree elt right_tree
388 in
389 aux 0 (pred (Array.length elts))
390 end
Something went wrong with that request. Please try again.