Not logged in.  Login/Logout/Register | List snippets | | Create snippet | Upload image | Upload data

429
LINES

< > BotCompany Repo | #1032621 // MegaCompactTreeSet [with more compact leaves than UltraCompactTreeSet, OK, 20 bytes per entry]

JavaX fragment (include) [tags: use-pretranspiled]

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  
}

Author comment

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: 115 / 258
Version history: 21 change(s)
Referenced in: [show references]