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

330
LINES

< > BotCompany Repo | #1030409 // IntPairTreeSet [derived from LongTreeSet]

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

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  
}

Author comment

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: 138 / 340
Version history: 2 change(s)
Referenced in: [show references]