1 | // based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html |
2 | // TODO: extends AbstractSet<Long> |
3 | sclass LongTreeSet implements Iterable<Long> { |
4 | // Left-leaning red-black BST with long values. |
5 | // This is the 2-3 version. |
6 | |
7 | private static final boolean RED = true; |
8 | private static final boolean BLACK = false; |
9 | |
10 | private Node root; // root of the BST |
11 | int size; // size of tree set |
12 | |
13 | // BST helper node data type |
14 | private sclass Node { |
15 | private long val; // associated data |
16 | private Node left, right; // links to left and right subtrees |
17 | private bool color; // color of parent link |
18 | |
19 | public Node(long val, boolean color) { |
20 | this.val = val; |
21 | this.color = color; |
22 | } |
23 | } |
24 | |
25 | *() {} |
26 | *(Cl<Long> cl) { addAll(cl); } |
27 | |
28 | // is node x red; false if x is null ? |
29 | private boolean isRed(Node x) { |
30 | if (x == null) return false; |
31 | return x.color == RED; |
32 | } |
33 | |
34 | public int size() { |
35 | ret size; |
36 | } |
37 | |
38 | public boolean isEmpty() { |
39 | return root == null; |
40 | } |
41 | |
42 | public bool add(long val) { |
43 | int oldSize = size; |
44 | root = put(root, val); |
45 | root.color = BLACK; |
46 | ifdef CompactTreeSet_debug assertTrue(check()); endifdef |
47 | ret size > oldSize; |
48 | } |
49 | |
50 | // insert the value in the subtree rooted at h |
51 | private Node put(Node h, long val) { |
52 | if (h == null) { ++size; ret new Node(val, RED); } |
53 | |
54 | int cmp = compare(val, h.val); |
55 | if (cmp < 0) h.left = put(h.left, val); |
56 | else if (cmp > 0) h.right = put(h.right, val); |
57 | else { /*h.val = val;*/ } // no overwriting |
58 | |
59 | // fix-up any right-leaning links |
60 | if (isRed(h.right) && !isRed(h.left)) h = rotateLeft(h); |
61 | if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h); |
62 | if (isRed(h.left) && isRed(h.right)) flipColors(h); |
63 | |
64 | ret h; |
65 | } |
66 | |
67 | // overridable if you want custom sorting |
68 | int compare(long a, long b) { |
69 | ret cmp(a, b); |
70 | } |
71 | |
72 | public bool remove(long key) { |
73 | if (!contains(key)) false; |
74 | |
75 | // if both children of root are black, set root to red |
76 | if (!isRed(root.left) && !isRed(root.right)) |
77 | root.color = RED; |
78 | |
79 | root = delete(root, key); |
80 | if (!isEmpty()) root.color = BLACK; |
81 | // assert check(); |
82 | true; |
83 | } |
84 | |
85 | // delete the key-value pair with the given key rooted at h |
86 | private Node delete(Node h, long key) { |
87 | // assert get(h, key) != null; |
88 | |
89 | if (compare(key, h.val) < 0) { |
90 | if (!isRed(h.left) && !isRed(h.left.left)) |
91 | h = moveRedLeft(h); |
92 | h.left = delete(h.left, key); |
93 | } |
94 | else { |
95 | if (isRed(h.left)) |
96 | h = rotateRight(h); |
97 | if (compare(key, h.val) == 0 && (h.right == null)) { |
98 | --size; null; |
99 | } if (!isRed(h.right) && !isRed(h.right.left)) |
100 | h = moveRedRight(h); |
101 | if (compare(key, h.val) == 0) { |
102 | --size; |
103 | Node x = min(h.right); |
104 | h.val = x.val; |
105 | // h.val = get(h.right, min(h.right).val); |
106 | // h.val = min(h.right).val; |
107 | h.right = deleteMin(h.right); |
108 | } |
109 | else h.right = delete(h.right, key); |
110 | } |
111 | return balance(h); |
112 | } |
113 | |
114 | // make a left-leaning link lean to the right |
115 | private Node rotateRight(Node h) { |
116 | // assert (h != null) && isRed(h.left); |
117 | Node x = h.left; |
118 | h.left = x.right; |
119 | x.right = h; |
120 | x.color = x.right.color; |
121 | x.right.color = RED; |
122 | return x; |
123 | } |
124 | |
125 | // make a right-leaning link lean to the left |
126 | private Node rotateLeft(Node h) { |
127 | // assert (h != null) && isRed(h.right); |
128 | Node x = h.right; |
129 | h.right = x.left; |
130 | x.left = h; |
131 | x.color = x.left.color; |
132 | x.left.color = RED; |
133 | return x; |
134 | } |
135 | |
136 | // flip the colors of a node and its two children |
137 | private void flipColors(Node h) { |
138 | // h must have opposite color of its two children |
139 | // assert (h != null) && (h.left != null) && (h.right != null); |
140 | // assert (!isRed(h) && isRed(h.left) && isRed(h.right)) |
141 | // || (isRed(h) && !isRed(h.left) && !isRed(h.right)); |
142 | h.color = !h.color; |
143 | h.left.color = !h.left.color; |
144 | h.right.color = !h.right.color; |
145 | } |
146 | |
147 | // Assuming that h is red and both h.left and h.left.left |
148 | // are black, make h.left or one of its children red. |
149 | private Node moveRedLeft(Node h) { |
150 | // assert (h != null); |
151 | // assert isRed(h) && !isRed(h.left) && !isRed(h.left.left); |
152 | |
153 | flipColors(h); |
154 | if (isRed(h.right.left)) { |
155 | h.right = rotateRight(h.right); |
156 | h = rotateLeft(h); |
157 | flipColors(h); |
158 | } |
159 | return h; |
160 | } |
161 | |
162 | // Assuming that h is red and both h.right and h.right.left |
163 | // are black, make h.right or one of its children red. |
164 | private Node moveRedRight(Node h) { |
165 | // assert (h != null); |
166 | // assert isRed(h) && !isRed(h.right) && !isRed(h.right.left); |
167 | flipColors(h); |
168 | if (isRed(h.left.left)) { |
169 | h = rotateRight(h); |
170 | flipColors(h); |
171 | } |
172 | return h; |
173 | } |
174 | |
175 | // restore red-black tree invariant |
176 | private Node balance(Node h) { |
177 | // assert (h != null); |
178 | |
179 | if (isRed(h.right)) h = rotateLeft(h); |
180 | if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h); |
181 | if (isRed(h.left) && isRed(h.right)) flipColors(h); |
182 | |
183 | return h; |
184 | } |
185 | |
186 | |
187 | /** |
188 | * Returns the height of the BST (for debugging). |
189 | * @return the height of the BST (a 1-node tree has height 0) |
190 | */ |
191 | public int height() { |
192 | return height(root); |
193 | } |
194 | |
195 | private int height(Node x) { |
196 | if (x == null) return -1; |
197 | return 1 + Math.max(height(x.left), height(x.right)); |
198 | } |
199 | |
200 | public bool contains(long val) { |
201 | ret find(root, val) != null; |
202 | } |
203 | |
204 | Node find(Node x, long key) { |
205 | while (x != null) { |
206 | int cmp = compare(key, x.val); |
207 | if (cmp < 0) x = x.left; |
208 | else if (cmp > 0) x = x.right; |
209 | else ret x; |
210 | } |
211 | null; |
212 | } |
213 | |
214 | private boolean check() { |
215 | if (!is23()) println("Not a 2-3 tree"); |
216 | if (!isBalanced()) println("Not balanced"); |
217 | return is23() && isBalanced(); |
218 | } |
219 | |
220 | // Does the tree have no red right links, and at most one (left) |
221 | // red links in a row on any path? |
222 | private boolean is23() { return is23(root); } |
223 | private boolean is23(Node x) { |
224 | if (x == null) return true; |
225 | if (isRed(x.right)) return false; |
226 | if (x != root && isRed(x) && isRed(x.left)) |
227 | return false; |
228 | return is23(x.left) && is23(x.right); |
229 | } |
230 | |
231 | // do all paths from root to leaf have same number of black edges? |
232 | private boolean isBalanced() { |
233 | int black = 0; // number of black links on path from root to min |
234 | Node x = root; |
235 | while (x != null) { |
236 | if (!isRed(x)) black++; |
237 | x = x.left; |
238 | } |
239 | return isBalanced(root, black); |
240 | } |
241 | |
242 | // does every path from the root to a leaf have the given number of black links? |
243 | private boolean isBalanced(Node x, int black) { |
244 | if (x == null) return black == 0; |
245 | if (!isRed(x)) black--; |
246 | return isBalanced(x.left, black) && isBalanced(x.right, black); |
247 | } |
248 | |
249 | public void clear() { root = null; size = 0; } |
250 | |
251 | // the smallest key in subtree rooted at x; null if no such key |
252 | private Node min(Node x) { |
253 | // assert x != null; |
254 | while (x.left != null) x = x.left; |
255 | ret x; |
256 | } |
257 | |
258 | private Node deleteMin(Node h) { |
259 | if (h.left == null) |
260 | return null; |
261 | |
262 | if (!isRed(h.left) && !isRed(h.left.left)) |
263 | h = moveRedLeft(h); |
264 | |
265 | h.left = deleteMin(h.left); |
266 | ret balance(h); |
267 | } |
268 | |
269 | public Iterator<Long> iterator() { |
270 | ret new MyIterator; |
271 | } |
272 | |
273 | class MyIterator extends ItIt<Long> { |
274 | new L<Node> path; |
275 | |
276 | *() { |
277 | fetch(root); |
278 | } |
279 | |
280 | void fetch(Node node) { |
281 | while (node != null) { |
282 | path.add(node); |
283 | node = node.left; |
284 | } |
285 | } |
286 | |
287 | public bool hasNext() { ret !path.isEmpty(); } |
288 | |
289 | public Long next() { |
290 | if (path.isEmpty()) fail("no more elements"); |
291 | Node node = popLast(path); |
292 | // last node is always a leaf, so left is null |
293 | // so proceed to fetch right branch |
294 | fetch(node.right); |
295 | ret node.val; |
296 | } |
297 | } |
298 | |
299 | // Returns the smallest key in the symbol table greater than or equal to {@code key}. |
300 | public Long ceiling(long key) { |
301 | Node x = ceiling(root, key); |
302 | ret x == null ? null : x.val; |
303 | } |
304 | |
305 | // the smallest key in the subtree rooted at x greater than or equal to the given key |
306 | Node ceiling(Node x, long key) { |
307 | if (x == null) null; |
308 | int cmp = compare(key, x.val); |
309 | if (cmp == 0) ret x; |
310 | if (cmp > 0) ret ceiling(x.right, key); |
311 | Node t = ceiling(x.left, key); |
312 | if (t != null) ret t; |
313 | else ret x; |
314 | } |
315 | |
316 | public Long floor(long key) { |
317 | Node x = floor(root, key); |
318 | ret x == null ? null : x.val; |
319 | } |
320 | |
321 | // the largest key in the subtree rooted at x less than or equal to the given key |
322 | Node floor(Node x, long key) { |
323 | if (x == null) null; |
324 | int cmp = compare(key, x.val); |
325 | if (cmp == 0) ret x; |
326 | if (cmp < 0) ret floor(x.left, key); |
327 | Node t = floor(x.right, key); |
328 | if (t != null) ret t; |
329 | else ret x; |
330 | } |
331 | |
332 | void addAll(Cl<Long> l) { |
333 | fOr (long x : l) add(x); |
334 | } |
335 | } |
Began life as a copy of #1030399
download show line numbers debug dex old transpilations
Travelled to 4 computer(s): bhatertpkbcr, mqqgnosmbjvj, pyentgdyhuwx, vouqrxazstgt
No comments. add comment
Snippet ID: | #1030407 |
Snippet name: | LongTreeSet [backup before AbstractSet] |
Eternal ID of this version: | #1030407/1 |
Text MD5: | f9599b5fe7e6c073d04fca39170136c9 |
Author: | stefan |
Category: | javax |
Type: | JavaX fragment (include) |
Public (visible to everyone): | Yes |
Archived (hidden from active list): | No |
Created/modified: | 2020-12-13 17:19:14 |
Source code size: | 9465 bytes / 335 lines |
Pitched / IR pitched: | No / No |
Views / Downloads: | 169 / 192 |
Referenced in: | [show references] |