LongTreeSet [backup before AbstractSet] [1030407]

// based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html
// TODO: extends AbstractSet<Long>
sclass LongTreeSet implements Iterable<Long> {
  // Left-leaning red-black BST with long values.
  // This is the 2-3 version.
 
  private static final boolean RED   = true;
  private static final boolean BLACK = false;

  private Node root;     // root of the BST
  int size; // size of tree set

  // BST helper node data type
  private sclass Node {
    private long val;          // associated data
    private Node left, right;  // links to left and right subtrees
    private bool color;        // color of parent link

    public Node(long val, boolean color) {
      this.val = val;
      this.color = color;
    }
  }
  
  *() {}
  *(Cl<Long> cl) { addAll(cl); }

  // is node x red; false if x is null ?
  private boolean isRed(Node x) {
      if (x == null) return false;
      return x.color == RED;
  }

  public int size() {
    ret size;
  }

  public boolean isEmpty() {
      return root == null;
  }

  public bool add(long val) {
    int oldSize = size;
    root = put(root, val);
    root.color = BLACK;
    ifdef CompactTreeSet_debug assertTrue(check()); endifdef
    ret size > oldSize;
  }

  // insert the value in the subtree rooted at h
  private Node put(Node h, long val) { 
    if (h == null) { ++size; ret new Node(val, RED); }

    int cmp = compare(val, h.val);
    if      (cmp < 0) h.left  = put(h.left,  val); 
    else if (cmp > 0) h.right = put(h.right, val); 
    else              { /*h.val = val;*/ } // no overwriting

    // fix-up any right-leaning links
    if (isRed(h.right) && !isRed(h.left))      h = rotateLeft(h);
    if (isRed(h.left)  &&  isRed(h.left.left)) h = rotateRight(h);
    if (isRed(h.left)  &&  isRed(h.right))     flipColors(h);

    ret h;
  }
  
  // overridable if you want custom sorting
  int compare(long a, long b) {
    ret cmp(a, b);
  }

  public bool remove(long key) { 
    if (!contains(key)) false;

    // if both children of root are black, set root to red
    if (!isRed(root.left) && !isRed(root.right))
        root.color = RED;

    root = delete(root, key);
    if (!isEmpty()) root.color = BLACK;
    // assert check();
    true;
  }

  // delete the key-value pair with the given key rooted at h
  private Node delete(Node h, long key) { 
    // assert get(h, key) != null;

    if (compare(key, h.val) < 0)  {
      if (!isRed(h.left) && !isRed(h.left.left))
          h = moveRedLeft(h);
      h.left = delete(h.left, key);
    }
    else {
      if (isRed(h.left))
          h = rotateRight(h);
      if (compare(key, h.val) == 0 && (h.right == null)) {
          --size; null;
      } if (!isRed(h.right) && !isRed(h.right.left))
          h = moveRedRight(h);
      if (compare(key, h.val) == 0) {
          --size;
          Node x = min(h.right);
          h.val = x.val;
          // h.val = get(h.right, min(h.right).val);
          // h.val = min(h.right).val;
          h.right = deleteMin(h.right);
      }
      else h.right = delete(h.right, key);
    }
    return balance(h);
  }

  // make a left-leaning link lean to the right
  private Node rotateRight(Node h) {
      // assert (h != null) && isRed(h.left);
      Node x = h.left;
      h.left = x.right;
      x.right = h;
      x.color = x.right.color;
      x.right.color = RED;
      return x;
  }

  // make a right-leaning link lean to the left
  private Node rotateLeft(Node h) {
      // assert (h != null) && isRed(h.right);
      Node x = h.right;
      h.right = x.left;
      x.left = h;
      x.color = x.left.color;
      x.left.color = RED;
      return x;
  }

  // flip the colors of a node and its two children
  private void flipColors(Node h) {
      // h must have opposite color of its two children
      // assert (h != null) && (h.left != null) && (h.right != null);
      // assert (!isRed(h) &&  isRed(h.left) &&  isRed(h.right))
      //    || (isRed(h)  && !isRed(h.left) && !isRed(h.right));
      h.color = !h.color;
      h.left.color = !h.left.color;
      h.right.color = !h.right.color;
  }

  // Assuming that h is red and both h.left and h.left.left
  // are black, make h.left or one of its children red.
  private Node moveRedLeft(Node h) {
      // assert (h != null);
      // assert isRed(h) && !isRed(h.left) && !isRed(h.left.left);

      flipColors(h);
      if (isRed(h.right.left)) { 
          h.right = rotateRight(h.right);
          h = rotateLeft(h);
          flipColors(h);
      }
      return h;
  }

  // Assuming that h is red and both h.right and h.right.left
  // are black, make h.right or one of its children red.
  private Node moveRedRight(Node h) {
      // assert (h != null);
      // assert isRed(h) && !isRed(h.right) && !isRed(h.right.left);
      flipColors(h);
      if (isRed(h.left.left)) { 
          h = rotateRight(h);
          flipColors(h);
      }
      return h;
  }

  // restore red-black tree invariant
  private Node balance(Node h) {
      // assert (h != null);

      if (isRed(h.right))                      h = rotateLeft(h);
      if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);
      if (isRed(h.left) && isRed(h.right))     flipColors(h);

      return h;
  }


  /**
   * Returns the height of the BST (for debugging).
   * @return the height of the BST (a 1-node tree has height 0)
   */
  public int height() {
      return height(root);
  }
  
  private int height(Node x) {
      if (x == null) return -1;
      return 1 + Math.max(height(x.left), height(x.right));
  }
  
  public bool contains(long val) {
    ret find(root, val) != null;
  }
  
  Node find(Node x, long key) {
    while (x != null) {
      int cmp = compare(key, x.val);
      if      (cmp < 0) x = x.left;
      else if (cmp > 0) x = x.right;
      else              ret x;
    }
    null;
  }

  private boolean check() {
      if (!is23())             println("Not a 2-3 tree");
      if (!isBalanced())       println("Not balanced");
      return is23() && isBalanced();
  }

  // Does the tree have no red right links, and at most one (left)
  // red links in a row on any path?
  private boolean is23() { return is23(root); }
  private boolean is23(Node x) {
      if (x == null) return true;
      if (isRed(x.right)) return false;
      if (x != root && isRed(x) && isRed(x.left))
          return false;
      return is23(x.left) && is23(x.right);
  } 

  // do all paths from root to leaf have same number of black edges?
  private boolean isBalanced() { 
      int black = 0;     // number of black links on path from root to min
      Node x = root;
      while (x != null) {
          if (!isRed(x)) black++;
          x = x.left;
      }
      return isBalanced(root, black);
  }

  // does every path from the root to a leaf have the given number of black links?
  private boolean isBalanced(Node x, int black) {
      if (x == null) return black == 0;
      if (!isRed(x)) black--;
      return isBalanced(x.left, black) && isBalanced(x.right, black);
  }
  
  public void clear() { root = null; size = 0; }
  
  // the smallest key in subtree rooted at x; null if no such key
  private Node min(Node x) { 
    // assert x != null;
    while (x.left != null) x = x.left;
    ret x;
  }
  
  private Node deleteMin(Node h) { 
    if (h.left == null)
      return null;

    if (!isRed(h.left) && !isRed(h.left.left))
      h = moveRedLeft(h);

    h.left = deleteMin(h.left);
    ret balance(h);
  }
  
  public Iterator<Long> iterator() {
    ret new MyIterator;
  }

  class MyIterator extends ItIt<Long> {
    new L<Node> path;
    
    *() {
      fetch(root);
    }
    
    void fetch(Node node) {
      while (node != null) {
        path.add(node);
        node = node.left;
      }
    }
    
    public bool hasNext() { ret !path.isEmpty(); }

    public Long next() {
      if (path.isEmpty()) fail("no more elements");
      Node node = popLast(path);
      // last node is always a leaf, so left is null
      // so proceed to fetch right branch
      fetch(node.right);
      ret node.val;
    }
  }
  
  // Returns the smallest key in the symbol table greater than or equal to {@code key}.
  public Long ceiling(long key) {
    Node x = ceiling(root, key);
    ret x == null ? null : x.val;
  }

  // the smallest key in the subtree rooted at x greater than or equal to the given key
  Node ceiling(Node x, long key) {  
    if (x == null) null;
    int cmp = compare(key, x.val);
    if (cmp == 0) ret x;
    if (cmp > 0)  ret ceiling(x.right, key);
    Node t = ceiling(x.left, key);
    if (t != null) ret t; 
    else           ret x;
  }
  
  public Long floor(long key) {
    Node x = floor(root, key);
    ret x == null ? null : x.val;
  }    

  // the largest key in the subtree rooted at x less than or equal to the given key
  Node floor(Node x, long key) {
    if (x == null) null;
    int cmp = compare(key, x.val);
    if (cmp == 0) ret x;
    if (cmp < 0)  ret floor(x.left, key);
    Node t = floor(x.right, key);
    if (t != null) ret t; 
    else           ret x;
  }
  
  void addAll(Cl<Long> l) {
    fOr (long x : l) add(x);
  }
}

Travelled to 4 computer(s): bhatertpkbcr, mqqgnosmbjvj, pyentgdyhuwx, vouqrxazstgt

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:	112 / 132
Referenced in:	[show references]

< > BotCompany Repo | #1030407 // LongTreeSet [backup before AbstractSet]

JavaX fragment (include)

Author comment

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	}