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