Libraryless. Click here for Pure Java version (4390L/26K).
1 | // based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html |
2 | // TODO: implement NavigableSet |
3 | sclass MegaCompactTreeSet<A> extends AbstractSet<A> { |
4 | |
5 | // A symbol table implemented using a left-leaning red-black BST. |
6 | // This is the 2-3 version. |
7 | |
8 | private static final boolean RED = true; |
9 | private static final boolean BLACK = false; |
10 | |
11 | private Node<A> root; // root of the BST |
12 | int size; // size of tree set |
13 | |
14 | // BST helper node data type |
15 | abstract sclass Node<A> { |
16 | A val; // associated data |
17 | |
18 | Node<A> left() { null; } // get left subtree |
19 | abstract Node<A> setLeft(Node<A> left); // set left subtree - return potentially replaced node |
20 | |
21 | Node<A> right() { null; } // get right subtree |
22 | abstract Node<A> setRight(Node<A> right); // set right subtree - return potentially replaced node |
23 | |
24 | abstract bool color(); |
25 | abstract Node<A> convertToBlack(); |
26 | abstract Node<A> convertToRed(); |
27 | abstract Node<A> invertColor(); |
28 | Node<A> convertToColor(bool color) { ret color == RED ? convertToRed() : convertToBlack(); } |
29 | |
30 | bool isLeaf() { ret left() == null && right() == null; } |
31 | } |
32 | |
33 | asclass NonLeaf<A> extends Node<A> { |
34 | Node<A> left, right; |
35 | |
36 | Node<A> left() { ret left; } |
37 | Node<A> setLeft(Node<A> left) { |
38 | this.left = left; |
39 | if (left == null && right() == null) ret newLeaf(color(), val); |
40 | this; |
41 | } |
42 | |
43 | Node<A> right() { ret right; } |
44 | Node<A> setRight(Node<A> right) { |
45 | this.right = right; |
46 | if (right == null && left() == null) ret newLeaf(color(), val); |
47 | this; |
48 | } |
49 | } |
50 | |
51 | sclass BlackNode<A> extends NonLeaf<A> { |
52 | *(A *val, Node<A> *left, Node<A> *right) {} |
53 | |
54 | bool color() { ret BLACK; } |
55 | Node<A> convertToBlack() { this; } |
56 | Node<A> convertToRed() { ret new RedNode(val, left, right); } |
57 | Node invertColor() { ret convertToRed(); } |
58 | } |
59 | |
60 | sclass RedNode<A> extends NonLeaf<A> { |
61 | *(A *val, Node<A> *left, Node<A> *right) {} |
62 | |
63 | bool color() { ret RED; } |
64 | Node<A> convertToBlack() { ret new BlackNode(val, left, right); } |
65 | Node<A> convertToRed() { this; } |
66 | Node invertColor() { ret convertToBlack(); } |
67 | } |
68 | |
69 | sclass BlackLeaf<A> extends Node<A> { |
70 | *(A *val) {} |
71 | |
72 | Node<A> setLeft(Node<A> left) { |
73 | ret new BlackNode(val, left, null); |
74 | } |
75 | |
76 | Node<A> setRight(Node<A> right) { |
77 | ret new BlackNode(val, null, right); |
78 | } |
79 | |
80 | bool color() { ret BLACK; } |
81 | Node<A> convertToBlack() { this; } |
82 | Node<A> convertToRed() { ret new RedLeaf(val); } |
83 | Node invertColor() { ret convertToRed(); } |
84 | } |
85 | |
86 | sclass RedLeaf<A> extends Node<A> { |
87 | *(A *val) {} |
88 | |
89 | Node<A> setLeft(Node<A> left) { |
90 | ret new RedNode(val, left, null); |
91 | } |
92 | |
93 | Node<A> setRight(Node<A> right) { |
94 | ret new RedNode(val, null, right); |
95 | } |
96 | |
97 | bool color() { ret RED; } |
98 | Node<A> convertToBlack() { ret new BlackLeaf(val); } |
99 | Node<A> convertToRed() { this; } |
100 | Node invertColor() { ret convertToBlack(); } |
101 | } |
102 | |
103 | *() {} |
104 | *(Cl<? extends A> cl) { addAll(cl); } |
105 | |
106 | // returns false on null (algorithm needs this) |
107 | static bool <A> isRed(Node<A> x) { |
108 | ret x != null && x.color() == RED; |
109 | } |
110 | |
111 | static <A> Node<A> newLeaf(bool color, A val) { |
112 | ret color == RED ? new RedLeaf(val) : new BlackLeaf(val); |
113 | } |
114 | |
115 | public int size() { |
116 | ret size; |
117 | } |
118 | |
119 | public bool isEmpty() { |
120 | ret root == null; |
121 | } |
122 | |
123 | public bool add(A val) { |
124 | int oldSize = size; |
125 | root = put(root, val); |
126 | root = root.convertToBlack(); |
127 | ifdef CompactTreeSet_debug assertTrue(check()); endifdef |
128 | ret size > oldSize; |
129 | } |
130 | |
131 | // insert the value in the subtree rooted at h |
132 | private Node<A> put(Node<A> h, A val) { |
133 | if (h == null) { ++size; ret new RedLeaf(val); } |
134 | |
135 | int cmp = compare(val, h.val); |
136 | if (cmp < 0) h = h.setLeft(put(h.left(), val)); |
137 | else if (cmp > 0) h = h.setRight(put(h.right(), val)) ; |
138 | else { /*h.val = val;*/ } // no overwriting |
139 | |
140 | // fix-up any right-leaning links |
141 | if (isRed(h.right()) && !isRed(h.left())) h = rotateLeft(h); |
142 | if (isRed(h.left()) && isRed(h.left().left())) h = rotateRight(h); |
143 | if (isRed(h.left()) && isRed(h.right())) h = flipColors(h); |
144 | |
145 | ret h; |
146 | } |
147 | |
148 | // override me if you wish |
149 | int compare(A a, A b) { |
150 | ret cmp(a, b); |
151 | } |
152 | |
153 | public bool remove(O key) { |
154 | if (!contains(key)) false; |
155 | |
156 | // if both children of root are black, set root to red |
157 | if (!isRed(root.left()) && !isRed(root.right())) |
158 | root = root.convertToRed(); |
159 | |
160 | root = delete(root, (A) key); |
161 | if (!isEmpty()) root = root.convertToBlack(); |
162 | // assert check(); |
163 | true; |
164 | } |
165 | |
166 | // delete the key-value pair with the given key rooted at h |
167 | private Node delete(Node<A> h, A key) { |
168 | // assert get(h, key) != null; |
169 | |
170 | if (compare(key, h.val) < 0) { |
171 | if (!isRed(h.left()) && !isRed(h.left().left())) |
172 | h = moveRedLeft(h); |
173 | h = h.setLeft(delete(h.left(), key)); |
174 | } |
175 | else { |
176 | if (isRed(h.left())) |
177 | h = rotateRight(h); |
178 | if (compare(key, h.val) == 0 && (h.right() == null)) { |
179 | --size; null; |
180 | } if (!isRed(h.right()) && !isRed(h.right().left())) |
181 | h = moveRedRight(h); |
182 | if (compare(key, h.val) == 0) { |
183 | --size; |
184 | Node<A> x = min(h.right()); |
185 | h.val = x.val; |
186 | // h.val = get(h.right(), min(h.right()).val); |
187 | // h.val = min(h.right()).val; |
188 | h = h.setRight(deleteMin(h.right())); |
189 | } |
190 | else h = h.setRight(delete(h.right(), key)); |
191 | } |
192 | return balance(h); |
193 | } |
194 | |
195 | // make a left-leaning link lean to the right |
196 | private Node<A> rotateRight(Node<A> h) { |
197 | // assert (h != null) && isRed(h.left()); |
198 | Node<A> x = h.left(); |
199 | h = h.setLeft(x.right()); |
200 | x = x.setRight(h); |
201 | x = x.convertToColor(x.right().color()); |
202 | x = x.setRight(x.right().convertToRed()); |
203 | ret x; |
204 | } |
205 | |
206 | // make a right-leaning link lean to the left |
207 | private Node<A> rotateLeft(Node<A> h) { |
208 | // assert (h != null) && isRed(h.right()); |
209 | Node<A> x = h.right(); |
210 | h = h.setRight(x.left()); |
211 | x = x.setLeft(h); |
212 | x = x.convertToColor(x.left().color()); |
213 | x = x.setLeft(x.left().convertToRed()); |
214 | ret x; |
215 | } |
216 | |
217 | // flip the colors of a node and its two children |
218 | private Node<A> flipColors(Node<A> h) { |
219 | // h must have opposite color of its two children |
220 | // assert (h != null) && (h.left() != null) && (h.right() != null); |
221 | // assert (!isRed(h) && isRed(h.left()) && isRed(h.right())) |
222 | // || (isRed(h) && !isRed(h.left()) && !isRed(h.right())); |
223 | h = h.setLeft(h.left().invertColor()); |
224 | h = h.setRight(h.right().invertColor()); |
225 | ret h.invertColor(); |
226 | } |
227 | |
228 | // Assuming that h is red and both h.left() and h.left().left() |
229 | // are black, make h.left() or one of its children red. |
230 | private Node<A> moveRedLeft(Node<A> h) { |
231 | // assert (h != null); |
232 | // assert isRed(h) && !isRed(h.left()) && !isRed(h.left().left()); |
233 | |
234 | h = flipColors(h); |
235 | if (isRed(h.right().left())) { |
236 | h = h.setRight(rotateRight(h.right())); |
237 | h = rotateLeft(h); |
238 | h = flipColors(h); |
239 | } |
240 | ret h; |
241 | } |
242 | |
243 | // Assuming that h is red and both h.right() and h.right().left() |
244 | // are black, make h.right() or one of its children red. |
245 | private Node<A> moveRedRight(Node<A> h) { |
246 | // assert (h != null); |
247 | // assert isRed(h) && !isRed(h.right()) && !isRed(h.right().left()); |
248 | h = flipColors(h); |
249 | if (isRed(h.left().left())) { |
250 | h = rotateRight(h); |
251 | h = flipColors(h); |
252 | } |
253 | ret h; |
254 | } |
255 | |
256 | // restore red-black tree invariant |
257 | private Node<A> balance(Node<A> h) { |
258 | // assert (h != null); |
259 | |
260 | if (isRed(h.right())) h = rotateLeft(h); |
261 | if (isRed(h.left()) && isRed(h.left().left())) h = rotateRight(h); |
262 | if (isRed(h.left()) && isRed(h.right())) h = flipColors(h); |
263 | |
264 | ret h; |
265 | } |
266 | |
267 | |
268 | /** |
269 | * Returns the height of the BST (for debugging). |
270 | * @return the height of the BST (a 1-node tree has height 0) |
271 | */ |
272 | public int height() { |
273 | ret height(root); |
274 | } |
275 | |
276 | private int height(Node<A> x) { |
277 | if (x == null) return -1; |
278 | return 1 + Math.max(height(x.left()), height(x.right())); |
279 | } |
280 | |
281 | public bool contains(O val) { |
282 | ret find(root, (A) val) != null; |
283 | } |
284 | |
285 | public A find(A probeVal) { |
286 | Node<A> n = find(root, probeVal); |
287 | ret n == null ? null : n.val; |
288 | } |
289 | |
290 | // value associated with the given key in subtree rooted at x; null if no such key |
291 | private A get(Node<A> x, A key) { |
292 | x = find(x, key); |
293 | ret x == null ? null : x.val; |
294 | } |
295 | |
296 | Node<A> find(Node<A> x, A key) { |
297 | while (x != null) { |
298 | int cmp = compare(key, x.val); |
299 | if (cmp < 0) x = x.left(); |
300 | else if (cmp > 0) x = x.right(); |
301 | else ret x; |
302 | } |
303 | null; |
304 | } |
305 | |
306 | private boolean check() { |
307 | if (!is23()) println("Not a 2-3 tree"); |
308 | if (!isBalanced()) println("Not balanced"); |
309 | return is23() && isBalanced(); |
310 | } |
311 | |
312 | // Does the tree have no red right links, and at most one (left) |
313 | // red links in a row on any path? |
314 | private boolean is23() { return is23(root); } |
315 | private boolean is23(Node<A> x) { |
316 | if (x == null) true; |
317 | if (isRed(x.right())) false; |
318 | if (x != root && isRed(x) && isRed(x.left())) false; |
319 | ret is23(x.left()) && is23(x.right()); |
320 | } |
321 | |
322 | // do all paths from root to leaf have same number of black edges? |
323 | private bool isBalanced() { |
324 | int black = 0; // number of black links on path from root to min |
325 | Node<A> x = root; |
326 | while (x != null) { |
327 | if (!isRed(x)) black++; |
328 | x = x.left(); |
329 | } |
330 | ret isBalanced(root, black); |
331 | } |
332 | |
333 | // does every path from the root to a leaf have the given number of black links? |
334 | private boolean isBalanced(Node<A> x, int black) { |
335 | if (x == null) return black == 0; |
336 | if (!isRed(x)) black--; |
337 | return isBalanced(x.left(), black) && isBalanced(x.right(), black); |
338 | } |
339 | |
340 | public void clear() { root = null; size = 0; } |
341 | |
342 | // the smallest key in subtree rooted at x; null if no such key |
343 | private Node<A> min(Node<A> x) { |
344 | // assert x != null; |
345 | while (x.left() != null) x = x.left(); |
346 | ret x; |
347 | } |
348 | |
349 | private Node<A> deleteMin(Node<A> h) { |
350 | if (h.left() == null) |
351 | return null; |
352 | |
353 | if (!isRed(h.left()) && !isRed(h.left().left())) |
354 | h = moveRedLeft(h); |
355 | |
356 | h = h.setLeft(deleteMin(h.left())); |
357 | ret balance(h); |
358 | } |
359 | |
360 | public Iterator<A> iterator() { |
361 | ret new MyIterator; |
362 | } |
363 | |
364 | class MyIterator extends ItIt<A> { |
365 | new L<Node<A>> path; |
366 | |
367 | *() { |
368 | fetch(root); |
369 | } |
370 | |
371 | void fetch(Node<A> node) { |
372 | while (node != null) { |
373 | path.add(node); |
374 | node = node.left(); |
375 | } |
376 | } |
377 | |
378 | public bool hasNext() { ret !path.isEmpty(); } |
379 | |
380 | public A next() { |
381 | if (path.isEmpty()) fail("no more elements"); |
382 | Node<A> node = popLast(path); |
383 | // last node is always a leaf, so left is null |
384 | // so proceed to fetch right branch |
385 | fetch(node.right()); |
386 | ret node.val; |
387 | } |
388 | } |
389 | |
390 | // Returns the smallest key in the symbol table greater than or equal to {@code key}. |
391 | public A ceiling(A key) { |
392 | Node<A> x = ceiling(root, key); |
393 | ret x == null ? null : x.val; |
394 | } |
395 | |
396 | // the smallest key in the subtree rooted at x greater than or equal to the given key |
397 | Node<A> ceiling(Node<A> x, A key) { |
398 | if (x == null) null; |
399 | int cmp = compare(key, x.val); |
400 | if (cmp == 0) ret x; |
401 | if (cmp > 0) ret ceiling(x.right(), key); |
402 | Node<A> t = ceiling(x.left(), key); |
403 | if (t != null) ret t; |
404 | else ret x; |
405 | } |
406 | |
407 | public A floor(A key) { |
408 | Node<A> x = floor(root, key); |
409 | ret x == null ? null : x.val; |
410 | } |
411 | |
412 | // the largest key in the subtree rooted at x less than or equal to the given key |
413 | Node<A> floor(Node<A> x, A key) { |
414 | if (x == null) null; |
415 | int cmp = compare(key, x.val); |
416 | if (cmp == 0) ret x; |
417 | if (cmp < 0) ret floor(x.left(), key); |
418 | Node<A> t = floor(x.right(), key); |
419 | if (t != null) ret t; |
420 | else ret x; |
421 | } |
422 | |
423 | void testInternalStructure(Node node default root) { |
424 | if (node == null) ret; |
425 | assertTrue(className(node), !node.isLeaf() == node instanceof NonLeaf); |
426 | testInternalStructure(node.left()); |
427 | testInternalStructure(node.right()); |
428 | } |
429 | } |
Began life as a copy of #1030938
download show line numbers debug dex old transpilations
Travelled to 3 computer(s): bhatertpkbcr, mowyntqkapby, mqqgnosmbjvj
No comments. add comment
Snippet ID: | #1032621 |
Snippet name: | MegaCompactTreeSet [with more compact leaves than UltraCompactTreeSet, OK, 20 bytes per entry] |
Eternal ID of this version: | #1032621/22 |
Text MD5: | f941dc0d5966600be20e068169fb2c97 |
Transpilation MD5: | b68a7a2cbe0fce0536f444e0c9192405 |
Author: | stefan |
Category: | javax |
Type: | JavaX fragment (include) |
Public (visible to everyone): | Yes |
Archived (hidden from active list): | No |
Created/modified: | 2021-09-29 21:14:55 |
Source code size: | 12776 bytes / 429 lines |
Pitched / IR pitched: | No / No |
Views / Downloads: | 170 / 340 |
Version history: | 21 change(s) |
Referenced in: | [show references] |