LogNArray v3 [works at least with add, remove, get, set] [1024413]

// see: https://en.wikipedia.org/wiki/Order_statistic_tree
// based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html
final sclass LogNArray<Value> extends RandomAccessAbstractList<Value> {

 // A symbol table implemented using a left-leaning red-black BST.
 // This is the 2-3 version.
 
  private static final boolean RED   = true;
  private static final boolean BLACK = false;

  private Node<Value> root;     // root of the BST

  // BST helper node data type
  private sclass Node<Value> {
      private Value val;         // associated data
      private Node<Value> left, right;  // links to left and right subtrees
      private boolean color;     // color of parent link
      private int size;          // subtree count

      public Node(Value val, boolean color, int size) {
          this.val = val;
          this.color = color;
          this.size = size;
      }
  }

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

  // number of node in subtree rooted at x; 0 if x is null
  private int size(Node<Value> x) {
      if (x == null) return 0;
      return x.size;
  } 


  public int size() {
      return size(root);
  }

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


  public void add(int idx, Value val) {
    if (idx < 0 || idx > size()) throw new IndexOutOfBoundsException(idx + " / " + size());
    
    root = add(root, idx, val);
    root.color = BLACK;
    ifdef LogNArray_debug assertTrue(check()); endifdef
  }

  // insert the value pair in index k of subtree rooted at h
  private Node<Value> add(Node<Value> h, int k, Value val) { 
    if (h == null) ret new Node(val, RED, 1);

    int t = size(h.left);
    if (k < t)
      // replace / fit in left child
      h.left = add(h.left, k, val);
    else if (k == t) {
      // move value to right child, replace value
      h.right = add(h.right, 0, h.val);
      h.val = val;
    } else
      // replace / fit in right child
      h.right = add(h.right, k-t-1, val); 

    // 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);
    h.size = size(h.left) + size(h.right) + 1;
    
    ifdef LogNArray_debug
      print("LogNArray.add: Node " + h.val + " new size: " + h.size);
    endifdef

    ret h;
  }

  public Value remove(int idx) { 
    Value oldValue = get(idx);
    
    // if both children of root are black, set root to red
    if (!isRed(root.left) && !isRed(root.right))
      root.color = RED;

    root = remove(root, idx);
    if (root != null) root.color = BLACK;
    
    ifdef LogNArray_debug assertTrue(check()); endifdef
    
    ret oldValue;
  }

  private Node<Value> remove(Node<Value> h, int k) { 
    int t = size(h.left);
    if (k < t) { // remove from left child
      if (!isRed(h.left) && !isRed(h.left.left))
        h = moveRedLeft(h);
      h.left = remove(h.left, k);
    } else { // remove from center or right child
      if (isRed(h.left)) {
        h = rotateRight(h);
        t = size(h.left);
      }
      if (h.right == null) null; // drop whole node
      if (!isRed(h.right) && !isRed(h.right.left)) {
        h = moveRedRight(h);
        t = size(h.left);
      }
      if (k == t) { // replace center value with min of right child
        h.val = nodeAtIndex(h.right, 0).val;
        h.right = remove(h.right, 0);
      }
      else h.right = remove(h.right, k-t-1);
    }
    ret balance(h);
  }

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

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

  // flip the colors of a node and its two children
  private void flipColors(Node<Value> 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<Value> moveRedLeft(Node<Value> 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<Value> moveRedRight(Node<Value> 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<Value> balance(Node<Value> 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);

      h.size = size(h.left) + size(h.right) + 1;
      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<Value> x) {
      if (x == null) return -1;
      return 1 + Math.max(height(x.left), height(x.right));
  }
  
  public Value get(int k) {
    ret nodeAtIndex(k).val;
  }

  public Value set(int k, Value val) {
    Node<Value> n = nodeAtIndex(k);
    Value oldValue = n.val;
    n.val = val;
    ret oldValue;
  }

  public Node<Value> nodeAtIndex(int k) {
      if (k < 0 || k >= size()) {
          throw new IndexOutOfBoundsException(k + " / " + size());
      }
      ret nodeAtIndex(root, k);
  }

  // the key of rank k in the subtree rooted at x
  private Node<Value> nodeAtIndex(Node<Value> x, int k) {
      // assert x != null;
      // assert k >= 0 && k < size(x);
      int t = size(x.left); 
      if      (t > k) return nodeAtIndex(x.left,  k); 
      else if (t < k) return nodeAtIndex(x.right, k-t-1); 
      else            return x; 
  } 
  
  private boolean check() {
      if (!isSizeConsistent()) println("Subtree counts not consistent");
      if (!is23())             println("Not a 2-3 tree");
      if (!isBalanced())       println("Not balanced");
      return isSizeConsistent() && is23() && isBalanced();
  }

  // are the size fields correct?
  private boolean isSizeConsistent() { return isSizeConsistent(root); }
  private boolean isSizeConsistent(Node<Value> x) {
      if (x == null) return true;
      if (x.size != size(x.left) + size(x.right) + 1) return false;
      return isSizeConsistent(x.left) && isSizeConsistent(x.right);
  } 

  // 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<Value> 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<Value> 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<Value> 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; }
}

Travelled to 6 computer(s): bhatertpkbcr, mqqgnosmbjvj, pyentgdyhuwx, pzhvpgtvlbxg, tvejysmllsmz, vouqrxazstgt

Snippet ID:	#1024413
Snippet name:	LogNArray v3 [works at least with add, remove, get, set]
Eternal ID of this version:	#1024413/26
Text MD5:	6fc3403a5782213c637221c7068d6da8
Transpilation MD5:	23c93820767d2abf8afc106f3a33fd18
Author:	stefan
Category:	javax
Type:	JavaX fragment (include)
Public (visible to everyone):	Yes
Archived (hidden from active list):	No
Created/modified:	2019-08-12 14:23:09
Source code size:	9091 bytes / 290 lines
Pitched / IR pitched:	No / No
Views / Downloads:	540 / 1024
Version history:	25 change(s)
Referenced in:	[show references]

< > BotCompany Repo | #1024413 // LogNArray v3 [works at least with add, remove, get, set]

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

Author comment

1	// see: https://en.wikipedia.org/wiki/Order_statistic_tree
2	// based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html
3	final sclass LogNArray<Value> extends RandomAccessAbstractList<Value> {
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<Value> root; // root of the BST
12
13	// BST helper node data type
14	private sclass Node<Value> {
15	private Value val; // associated data
16	private Node<Value> left, right; // links to left and right subtrees
17	private boolean color; // color of parent link
18	private int size; // subtree count
19
20	public Node(Value val, boolean color, int size) {
21	this.val = val;
22	this.color = color;
23	this.size = size;
24	}
25	}
26
27	// is node x red; false if x is null ?
28	private boolean isRed(Node<Value> x) {
29	if (x == null) return false;
30	return x.color == RED;
31	}
32
33	// number of node in subtree rooted at x; 0 if x is null
34	private int size(Node<Value> x) {
35	if (x == null) return 0;
36	return x.size;
37	}
38
39
40	public int size() {
41	return size(root);
42	}
43
44	public boolean isEmpty() {
45	return root == null;
46	}
47
48
49	public void add(int idx, Value val) {
50	if (idx < 0 \|\| idx > size()) throw new IndexOutOfBoundsException(idx + " / " + size());
51
52	root = add(root, idx, val);
53	root.color = BLACK;
54	ifdef LogNArray_debug assertTrue(check()); endifdef
55	}
56
57	// insert the value pair in index k of subtree rooted at h
58	private Node<Value> add(Node<Value> h, int k, Value val) {
59	if (h == null) ret new Node(val, RED, 1);
60
61	int t = size(h.left);
62	if (k < t)
63	// replace / fit in left child
64	h.left = add(h.left, k, val);
65	else if (k == t) {
66	// move value to right child, replace value
67	h.right = add(h.right, 0, h.val);
68	h.val = val;
69	} else
70	// replace / fit in right child
71	h.right = add(h.right, k-t-1, val);
72
73	// fix-up any right-leaning links
74	if (isRed(h.right) && !isRed(h.left)) h = rotateLeft(h);
75	if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);
76	if (isRed(h.left) && isRed(h.right)) flipColors(h);
77	h.size = size(h.left) + size(h.right) + 1;
78
79	ifdef LogNArray_debug
80	print("LogNArray.add: Node " + h.val + " new size: " + h.size);
81	endifdef
82
83	ret h;
84	}
85
86	public Value remove(int idx) {
87	Value oldValue = get(idx);
88
89	// if both children of root are black, set root to red
90	if (!isRed(root.left) && !isRed(root.right))
91	root.color = RED;
92
93	root = remove(root, idx);
94	if (root != null) root.color = BLACK;
95
96	ifdef LogNArray_debug assertTrue(check()); endifdef
97
98	ret oldValue;
99	}
100
101	private Node<Value> remove(Node<Value> h, int k) {
102	int t = size(h.left);
103	if (k < t) { // remove from left child
104	if (!isRed(h.left) && !isRed(h.left.left))
105	h = moveRedLeft(h);
106	h.left = remove(h.left, k);
107	} else { // remove from center or right child
108	if (isRed(h.left)) {
109	h = rotateRight(h);
110	t = size(h.left);
111	}
112	if (h.right == null) null; // drop whole node
113	if (!isRed(h.right) && !isRed(h.right.left)) {
114	h = moveRedRight(h);
115	t = size(h.left);
116	}
117	if (k == t) { // replace center value with min of right child
118	h.val = nodeAtIndex(h.right, 0).val;
119	h.right = remove(h.right, 0);
120	}
121	else h.right = remove(h.right, k-t-1);
122	}
123	ret balance(h);
124	}
125
126	// make a left-leaning link lean to the right
127	private Node<Value> rotateRight(Node<Value> h) {
128	// assert (h != null) && isRed(h.left);
129	Node<Value> x = h.left;
130	h.left = x.right;
131	x.right = h;
132	x.color = x.right.color;
133	x.right.color = RED;
134	x.size = h.size;
135	h.size = size(h.left) + size(h.right) + 1;
136	return x;
137	}
138
139	// make a right-leaning link lean to the left
140	private Node<Value> rotateLeft(Node<Value> h) {
141	// assert (h != null) && isRed(h.right);
142	Node<Value> x = h.right;
143	h.right = x.left;
144	x.left = h;
145	x.color = x.left.color;
146	x.left.color = RED;
147	x.size = h.size;
148	h.size = size(h.left) + size(h.right) + 1;
149	return x;
150	}
151
152	// flip the colors of a node and its two children
153	private void flipColors(Node<Value> h) {
154	// h must have opposite color of its two children
155	// assert (h != null) && (h.left != null) && (h.right != null);
156	// assert (!isRed(h) && isRed(h.left) && isRed(h.right))
157	// \|\| (isRed(h) && !isRed(h.left) && !isRed(h.right));
158	h.color = !h.color;
159	h.left.color = !h.left.color;
160	h.right.color = !h.right.color;
161	}
162
163	// Assuming that h is red and both h.left and h.left.left
164	// are black, make h.left or one of its children red.
165	private Node<Value> moveRedLeft(Node<Value> h) {
166	// assert (h != null);
167	// assert isRed(h) && !isRed(h.left) && !isRed(h.left.left);
168
169	flipColors(h);
170	if (isRed(h.right.left)) {
171	h.right = rotateRight(h.right);
172	h = rotateLeft(h);
173	flipColors(h);
174	}
175	return h;
176	}
177
178	// Assuming that h is red and both h.right and h.right.left
179	// are black, make h.right or one of its children red.
180	private Node<Value> moveRedRight(Node<Value> h) {
181	// assert (h != null);
182	// assert isRed(h) && !isRed(h.right) && !isRed(h.right.left);
183	flipColors(h);
184	if (isRed(h.left.left)) {
185	h = rotateRight(h);
186	flipColors(h);
187	}
188	return h;
189	}
190
191	// restore red-black tree invariant
192	private Node<Value> balance(Node<Value> h) {
193	// assert (h != null);
194
195	if (isRed(h.right)) h = rotateLeft(h);
196	if (isRed(h.left) && isRed(h.left.left)) h = rotateRight(h);
197	if (isRed(h.left) && isRed(h.right)) flipColors(h);
198
199	h.size = size(h.left) + size(h.right) + 1;
200	return h;
201	}
202
203
204	/**
205	* Returns the height of the BST (for debugging).
206	* @return the height of the BST (a 1-node tree has height 0)
207	*/
208	public int height() {
209	return height(root);
210	}
211
212	private int height(Node<Value> x) {
213	if (x == null) return -1;
214	return 1 + Math.max(height(x.left), height(x.right));
215	}
216
217	public Value get(int k) {
218	ret nodeAtIndex(k).val;
219	}
220
221	public Value set(int k, Value val) {
222	Node<Value> n = nodeAtIndex(k);
223	Value oldValue = n.val;
224	n.val = val;
225	ret oldValue;
226	}
227
228	public Node<Value> nodeAtIndex(int k) {
229	if (k < 0 \|\| k >= size()) {
230	throw new IndexOutOfBoundsException(k + " / " + size());
231	}
232	ret nodeAtIndex(root, k);
233	}
234
235	// the key of rank k in the subtree rooted at x
236	private Node<Value> nodeAtIndex(Node<Value> x, int k) {
237	// assert x != null;
238	// assert k >= 0 && k < size(x);
239	int t = size(x.left);
240	if (t > k) return nodeAtIndex(x.left, k);
241	else if (t < k) return nodeAtIndex(x.right, k-t-1);
242	else return x;
243	}
244
245	private boolean check() {
246	if (!isSizeConsistent()) println("Subtree counts not consistent");
247	if (!is23()) println("Not a 2-3 tree");
248	if (!isBalanced()) println("Not balanced");
249	return isSizeConsistent() && is23() && isBalanced();
250	}
251
252	// are the size fields correct?
253	private boolean isSizeConsistent() { return isSizeConsistent(root); }
254	private boolean isSizeConsistent(Node<Value> x) {
255	if (x == null) return true;
256	if (x.size != size(x.left) + size(x.right) + 1) return false;
257	return isSizeConsistent(x.left) && isSizeConsistent(x.right);
258	}
259
260	// Does the tree have no red right links, and at most one (left)
261	// red links in a row on any path?
262	private boolean is23() { return is23(root); }
263	private boolean is23(Node<Value> x) {
264	if (x == null) return true;
265	if (isRed(x.right)) return false;
266	if (x != root && isRed(x) && isRed(x.left))
267	return false;
268	return is23(x.left) && is23(x.right);
269	}
270
271	// do all paths from root to leaf have same number of black edges?
272	private boolean isBalanced() {
273	int black = 0; // number of black links on path from root to min
274	Node<Value> x = root;
275	while (x != null) {
276	if (!isRed(x)) black++;
277	x = x.left;
278	}
279	return isBalanced(root, black);
280	}
281
282	// does every path from the root to a leaf have the given number of black links?
283	private boolean isBalanced(Node<Value> x, int black) {
284	if (x == null) return black == 0;
285	if (!isRed(x)) black--;
286	return isBalanced(x.left, black) && isBalanced(x.right, black);
287	}
288
289	public void clear() { root = null; }
290	}