1  import java.util.ArrayList;
  2  
  3  public class RBTree<E extends Comparable<E>> extends BST<E> {
  4    /** Create a default RB tree */
  5    public RBTree() {
  6    }
  7  
  8    /** Create an RB tree from an array of elements */
  9    public RBTree(E[] elements) {
 10      super(elements);
 11    }
 12  
 13    @Override /** Override createNewNode to create an RBTreeNode */
 14    protected RBTreeNode<E> createNewNode(E e) {
 15      return new RBTreeNode<E>(e);
 16    }
 17  
 18    @Override /** Override the insert method to 
 19      balance the tree if necessary */
 20    public boolean insert(E e) {
 21      boolean successful = super.insert(e);
 22      if (!successful)
 23        return false; // e is already in the tree
 24      else {
 25        ensureRBTree(e); 
 26      }
 27  
 28      return true; // e is inserted
 29    }
 30  
 31    /** Ensure that the tree is a red-black tree */
 32    private void ensureRBTree(E e) {
 33      // Get the path that leads to element e from the root 
 34      ArrayList<TreeNode<E>> path = path(e);
 35  
 36      int i = path.size() - 1; // Index to the current node in the path
 37      
 38      // u is the last node in the path. u contains element e
 39      RBTreeNode<E> u = (RBTreeNode<E>)(path.get(i));
 40          
 41      // v is the parent of of u, if exists
 42      RBTreeNode<E> v = (u == root) ? null : 
 43        (RBTreeNode<E>)(path.get(i - 1));
 44  
 45      u.setRed(); // It is OK to set u red    
 46            
 47      if (u == root) // If e is inserted as the root, set root black
 48        u.setBlack();
 49      else if (v.isRed()) 
 50        fixDoubleRed(u, v, path, i); // Fix double red violation at u
 51    }
 52    
 53    /** Fix double red violation at node u */
 54    private void fixDoubleRed(RBTreeNode<E> u, RBTreeNode<E> v, 
 55        ArrayList<TreeNode<E>> path, int i) {          
 56      // w is the grandparent of u
 57      RBTreeNode<E> w = (RBTreeNode<E>)(path.get(i - 2));
 58      RBTreeNode<E> parentOfw = (w == root) ? null : 
 59        (RBTreeNode<E>)path.get(i - 3);
 60  
 61      // Get v's sibling named x
 62      RBTreeNode<E> x = (w.left == v) ? 
 63        (RBTreeNode<E>)(w.right) : (RBTreeNode<E>)(w.left);
 64      
 65      if (x == null || x.isBlack()) {
 66        // Case 1: v's sibling x is black
 67        if (w.left == v && v.left == u) {
 68          // Case 1.1: u < v < w, Restructure and recolor nodes
 69          restructureRecolor(u, v, w, w, parentOfw);
 70      
 71          w.left = v.right; // v.right is y3 in Figure 11.7
 72          v.right = w;
 73        }
 74        else if (w.left == v && v.right == u) {
 75          // Case 1.2: v < u < w, Restructure and recolor nodes
 76          restructureRecolor(v, u, w, w, parentOfw);
 77          v.right = u.left;
 78          w.left = u.right;
 79          u.left = v;
 80          u.right = w;
 81        }
 82        else if (w.right == v && v.right == u) {
 83          // Case 1.3: w < v < u, Restructure and recolor nodes
 84          restructureRecolor(w, v, u, w, parentOfw);
 85          w.right = v.left;
 86          v.left = w;
 87        }
 88        else {
 89          // Case 1.4: w < u < v, Restructure and recolor nodes
 90          restructureRecolor(w, u, v, w, parentOfw);
 91          w.right = u.left;
 92          v.left = u.right;
 93          u.left = w;
 94          u.right = v;
 95        }
 96      }
 97      else { // Case 2: v's sibling x is red 
 98        // Recolor nodes
 99        w.setRed();
100        u.setRed();
101        ((RBTreeNode<E>)(w.left)).setBlack(); 
102        ((RBTreeNode<E>)(w.right)).setBlack();
103        
104        if (w == root) {
105          w.setBlack();     
106        }
107        else if (((RBTreeNode<E>)parentOfw).isRed()) {  
108          // Propagate along the path to fix new double red violation
109          u = w;
110          v = (RBTreeNode<E>)parentOfw;
111          fixDoubleRed(u, v, path, i - 2); // i – 2 propagates upward
112        }
113      }
114    }
115  
116    /** Connect b with parentOfw and recolor a, b, c for a < b < c */
117    private void restructureRecolor(RBTreeNode<E> a, RBTreeNode<E> b, 
118        RBTreeNode<E> c, RBTreeNode<E> w, RBTreeNode<E> parentOfw) {
119      if (parentOfw == null)
120        root = b;
121      else if (parentOfw.left == w)
122        parentOfw.left = b;
123      else
124        parentOfw.right = b;
125  
126      b.setBlack(); // b becomes the root in the subtree
127      a.setRed(); // a becomes the left child of b
128      c.setRed(); // c becomes the right child of b
129    }      
130    
131    @Override /** Delete an element from the RBTree.
132     * Return true if the element is deleted successfully
133     * Return false if the element is not in the tree */
134    public boolean delete(E e) {
135      // Locate the node to be deleted
136      TreeNode<E> current = root;
137      while (current != null) {
138        if (e.compareTo(current.element) < 0) {
139          current = current.left;
140        }
141        else if (e.compareTo(current.element) > 0) {
142          current = current.right;
143        }
144        else
145          break; // Element is in the tree pointed by current
146      }
147  
148      if (current == null)
149        return false; // Element is not in the tree
150  
151      java.util.ArrayList<TreeNode<E>> path;
152  
153      // current node is an internal node 
154      if (current.left != null && current.right != null) {
155        // Locate the rightmost node in the left subtree of current
156        TreeNode<E> rightMost = current.left;
157        while (rightMost.right != null) {
158          rightMost = rightMost.right; // Keep going to the right
159        }
160  
161        path = path(rightMost.element); // Get path before replacement
162  
163        // Replace the element in current by the element in rightMost
164        current.element = rightMost.element;
165      }
166      else
167        path = path(e); // Get path to current node
168  
169      // Delete the last node in the path and propagate if needed
170      deleteLastNodeInPath(path);
171      
172      size--; // After one element deleted
173      return true; // Element deleted
174    }
175  
176    /** Delete the last node from the path. */
177    public void deleteLastNodeInPath(ArrayList<TreeNode<E>> path) {
178      int i = path.size() - 1; // Index to the node in the path
179      // u is the last node in the path
180      RBTreeNode<E> u = (RBTreeNode<E>)(path.get(i));
181      RBTreeNode<E> parentOfu = (u == root) ? null :
182        (RBTreeNode<E>)(path.get(i - 1));
183      RBTreeNode<E> grandparentOfu = (parentOfu == null ||
184        parentOfu == root) ? null :
185        (RBTreeNode<E>)(path.get(i - 2));
186      RBTreeNode<E> childOfu = (u.left == null) ?
187        (RBTreeNode<E>)(u.right) : (RBTreeNode<E>)(u.left);
188  
189      // Delete node u. Connect childOfu with parentOfu
190      connectNewParent(parentOfu, u, childOfu);
191      
192      // Recolor the nodes and fix double black if needed
193      if (childOfu == root || u.isRed())
194        return; // Done if childOfu is root or if u is red 
195      else if (childOfu != null && childOfu.isRed()) 
196        childOfu.setBlack(); // Set it black, done
197      else // u is black, childOfu is null or black
198        // Fix double black on parentOfu
199        fixDoubleBlack(grandparentOfu, parentOfu, childOfu, path, i);
200    }
201  
202    /** Fix the double black problem at node parent */
203    private void fixDoubleBlack(
204        RBTreeNode<E> grandparent, RBTreeNode<E> parent, 
205        RBTreeNode<E> db, ArrayList<TreeNode<E>> path, int i) {
206      // Obtain y, y1, and y2
207      RBTreeNode<E> y = (parent.right == db) ? 
208        (RBTreeNode<E>)(parent.left) : (RBTreeNode<E>)(parent.right);
209      RBTreeNode<E> y1 = (RBTreeNode<E>)(y.left);
210      RBTreeNode<E> y2 = (RBTreeNode<E>)(y.right);
211  
212      if (y.isBlack() && y1 != null && y1.isRed()) {
213        if (parent.right == db) {
214          // Case 1.1: y is a left black sibling and y1 is red
215          connectNewParent(grandparent, parent, y);
216          recolor(parent, y, y1); // Adjust colors
217  
218          // Adjust child links
219          parent.left = y.right;
220          y.right = parent;
221        }
222        else {
223          // Case 1.3: y is a right black sibling and y1 is red        
224          connectNewParent(grandparent, parent, y1);
225          recolor(parent, y1, y); // Adjust colors
226  
227          // Adjust child links
228          parent.right = y1.left;
229          y.left = y1.right;
230          y1.left = parent;
231          y1.right = y;
232        }
233      }
234      else if (y.isBlack() && y2 != null && y2.isRed()) {
235        if (parent.right == db) {
236          // Case 1.2: y is a left black sibling and y2 is red
237          connectNewParent(grandparent, parent, y2);
238          recolor(parent, y2, y); // Adjust colors
239  
240          // Adjust child links
241          y.right = y2.left;
242          parent.left = y2.right;
243          y2.left = y;
244          y2.right = parent;
245        }
246        else {
247          // Case 1.4: y is a right black sibling and y2 is red        
248          connectNewParent(grandparent, parent, y);
249          recolor(parent, y, y2); // Adjust colors
250  
251          // Adjust child links
252          y.left = parent;
253          parent.right = y1;
254        }
255      }
256      else if (y.isBlack()) { 
257        // Case 2: y is black and y's children are black or null
258        y.setRed(); // Change y to red
259        if (parent.isRed())
260          parent.setBlack(); // Done
261        else if (parent != root) {
262          // Propagate double black to the parent node
263          // Fix new appearance of double black recursively
264          db = parent;
265          parent = grandparent;
266          grandparent = 
267            (i >= 3) ? (RBTreeNode<E>)(path.get(i - 3)) : null;
268          fixDoubleBlack(grandparent, parent, db, path, i - 1);
269        }
270      }
271      else { // y.isRed()
272        if (parent.right == db) {       
273          // Case 3.1: y is a left red child of parent
274          parent.left = y2;
275          y.right = parent;
276        }
277        else {
278          // Case 3.2: y is a right red child of parent
279          parent.right = y.left;
280          y.left = parent;
281        } 
282        
283        parent.setRed(); // Color parent red
284        y.setBlack(); // Color y black
285        connectNewParent(grandparent, parent, y); // y is new parent
286        fixDoubleBlack(y, parent, db, path, i - 1); 
287      }
288    }
289  
290    /** Recolor parent, newParent, and c. Case 1 removal */
291    private void recolor(RBTreeNode<E> parent, 
292        RBTreeNode<E> newParent, RBTreeNode<E> c) {
293      // Retain the parent's color for newParent
294      if (parent.isRed())
295        newParent.setRed(); 
296      else
297        newParent.setBlack();
298      
299      // c and parent become the children of newParent, set them black
300      parent.setBlack();  
301      c.setBlack();
302    }
303  
304    /** Connect newParent with grandParent */   
305    private void connectNewParent(RBTreeNode<E> grandparent,
306        RBTreeNode<E> parent, RBTreeNode<E> newParent) {
307      if (parent == root) {
308        root = newParent;
309        if (root != null)
310          newParent.setBlack();
311      }
312      else if (grandparent.left == parent)
313        grandparent.left = newParent;
314      else
315        grandparent.right = newParent;
316    }
317  
318    @Override /** Preorder traversal from a subtree */
319    protected void preorder(TreeNode<E> root) {
320      if (root == null) return;
321      System.out.print(root.element +
322        (((RBTreeNode<E>)root).isRed() ? " (red) " : " (black) "));
323      preorder(root.left);
324      preorder(root.right);
325    }
326  
327    /** RBTreeNode is TreeNode plus color indicator */
328    protected static class RBTreeNode<E extends Comparable<E>> extends
329        BST.TreeNode<E> {
330      private boolean red = true; // Indicate node color
331  
332      public RBTreeNode(E e) {
333        super(e);
334      }
335  
336      public boolean isRed() {
337        return red;
338      }
339  
340      public boolean isBlack() {
341        return!red;
342      }
343  
344      public void setBlack() {
345        red = false;
346      }
347  
348      public void setRed() {
349        red = true;
350      }
351  
352      int blackHeight;
353    }
354  }