1  import java.util.ArrayList;
  2  import java.util.List;
  3  
  4  /**
  5     This class implements a red-black tree whose
  6     nodes hold objects that implement the Comparable
  7     interface.
  8  */
  9  public class RedBlackTree
 10  {  
 11     Node root; // Package access, for testing
 12  
 13     static final int BLACK = 1;
 14     static final int RED = 0;
 15     private static final int NEGATIVE_RED = -1;
 16     private static final int DOUBLE_BLACK = 2;
 17  
 18     /**
 19        Constructs an empty tree.
 20     */
 21     public RedBlackTree()
 22     {  
 23        root = null;
 24     }
 25     
 26     /**
 27        Inserts a new node into the tree.
 28        @param obj the object to insert
 29     */
 30     public void add(Comparable obj) 
 31     {  
 32        Node newNode = new Node();
 33        newNode.data = obj;
 34        newNode.left = null;
 35        newNode.right = null;
 36        if (root == null) { root = newNode; }
 37        else { root.addNode(newNode); }
 38        fixAfterAdd(newNode);
 39     }
 40  
 41     /**
 42        Tries to find an object in the tree.
 43        @param obj the object to find
 44        @return true if the object is contained in the tree
 45     */
 46     public boolean find(Comparable obj)
 47     {
 48        Node current = root;
 49        while (current != null)
 50        {
 51           int d = current.data.compareTo(obj);
 52           if (d == 0) { return true; }
 53           else if (d > 0) { current = current.left; }
 54           else { current = current.right; }
 55        }
 56        return false;
 57     }
 58     
 59     /**
 60        Tries to remove an object from the tree. Does nothing
 61        if the object is not contained in the tree.
 62        @param obj the object to remove
 63     */
 64     public void remove(Comparable obj)
 65     {
 66        // Find node to be removed
 67  
 68        Node toBeRemoved = root;
 69        boolean found = false;
 70        while (!found && toBeRemoved != null)
 71        {
 72           int d = toBeRemoved.data.compareTo(obj);
 73           if (d == 0) { found = true; }
 74           else 
 75           {
 76              if (d > 0) { toBeRemoved = toBeRemoved.left; }
 77              else { toBeRemoved = toBeRemoved.right; }
 78           }
 79        }
 80  
 81        if (!found) { return; }
 82  
 83        // toBeRemoved contains obj
 84  
 85        // If one of the children is empty, use the other
 86  
 87        if (toBeRemoved.left == null || toBeRemoved.right == null)
 88        {
 89           Node newChild;
 90           if (toBeRemoved.left == null) { newChild = toBeRemoved.right; }
 91           else { newChild = toBeRemoved.left; }
 92  
 93           fixBeforeRemove(toBeRemoved); 
 94           replaceWith(toBeRemoved, newChild);
 95           return;
 96        }
 97        
 98        // Neither subtree is empty
 99  
100        // Find smallest element of the right subtree
101  
102        Node smallest = toBeRemoved.right;
103        while (smallest.left != null)
104        {
105           smallest = smallest.left;
106        }
107  
108        // smallest contains smallest child in right subtree
109           
110        // Move contents, unlink child
111  
112        toBeRemoved.data = smallest.data;
113        fixBeforeRemove(smallest);
114        replaceWith(smallest, smallest.right);
115     }
116     
117     /**
118        Yields the contents of the tree in sorted order
119        @return all data items traversed in inorder, with a space after each item
120     */
121     public String toString()
122     {  
123        return toString(root);
124     }  
125  
126  
127     /**
128        Yields the contents of the subtree with the given root in sorted order
129        @param parent the root of the subtree
130        @return all data items traversed in inorder, with a space after each item
131     */
132     private static String toString(Node parent)
133     {  
134        if (parent == null) { return ""; }
135        return toString(parent.left) + parent.data + " " + toString(parent.right);
136     }
137  
138     /**
139        A node of a red-black tree stores a data item and references
140        of the child nodes to the left and to the right. The color
141        is the "cost" of passing the node; 1 for black or 0 for red.
142        Temporarily, it may be set to 2 or -1. 
143     */
144     static class Node
145     {  
146        public Comparable data;
147        public Node left;
148        public Node right;
149        public Node parent;
150        public int color;
151        
152        /**
153           Constructs a red node with no data.
154        */
155        public Node() {}
156        
157        /**
158           Sets the left child and updates its parent reference.
159           @param child the new left child
160        */
161        public void setLeftChild(Node child)
162        {
163           left = child;
164           if (child != null) { child.parent = this; }
165        }
166        
167        /**
168           Sets the right child and updates its parent reference.
169           @param child the new right child
170        */
171        public void setRightChild(Node child)
172        {
173           right = child;
174           if (child != null) { child.parent = this; }
175        }
176        
177        /**
178           Inserts a new node as a descendant of this node.
179           @param newNode the node to insert
180        */
181        public void addNode(Node newNode)
182        {  
183           int comp = newNode.data.compareTo(data);
184           if (comp < 0)
185           {  
186              if (left == null) 
187              {
188                 left = newNode;
189                 left.parent = this;
190              }
191              else { left.addNode(newNode); }
192           }
193           else if (comp > 0)
194           {  
195              if (right == null) 
196              {
197                 right = newNode;
198                 right.parent = this;
199              }
200              else { right.addNode(newNode); }
201           }
202        }
203     }
204  
205     /**
206        Updates the parent's and replacement node's links when this node is replaced.
207        Also updates the root reference if it is replaced.
208        @param toBeReplaced the node that is to be replaced
209        @param replacement the node that replaces that node
210     */
211     private void replaceWith(Node toBeReplaced, Node replacement)
212     {
213        if (toBeReplaced.parent == null) 
214        { 
215           replacement.parent = null; 
216           root = replacement; 
217        }
218        else if (toBeReplaced == toBeReplaced.parent.left) 
219        { 
220           toBeReplaced.parent.setLeftChild(replacement); 
221        }
222        else 
223        { 
224           toBeReplaced.parent.setRightChild(replacement); 
225        }
226     }
227  
228     /**
229        Restores the tree to a red-black tree after a node has been added.
230        @param newNode the node that has been added
231     */
232     private void fixAfterAdd(Node newNode)
233     {
234        if (newNode.parent == null) 
235        { 
236           newNode.color = BLACK; 
237        }
238        else
239        {
240           newNode.color = RED;
241           if (newNode.parent.color == RED) { fixDoubleRed(newNode); }
242        }
243     }
244  
245     /** 	
246       Fixes the tree so that it is a red-black tree after a node has been removed.
247       @param toBeRemoved the node that is to be removed
248     */
249     private void fixBeforeRemove(Node toBeRemoved)
250     {
251        if (toBeRemoved.color == RED) { return; }
252  
253        if (toBeRemoved.left != null || toBeRemoved.right != null) // It is not a leaf
254        {
255           // Color the child black
256           if (toBeRemoved.left == null) { toBeRemoved.right.color = BLACK; }
257           else { toBeRemoved.left.color = BLACK; }
258        }	   
259        else { bubbleUp(toBeRemoved.parent); }
260     }
261     
262     /**
263        Move a charge from two children of a parent
264        @param parent a node with two children, or null (in which case nothing is done)
265     */
266     private void bubbleUp(Node parent)
267     {
268        if (parent == null) { return; }
269        parent.color++;
270        parent.left.color--;
271        parent.right.color--;
272  	   
273        if (bubbleUpFix(parent.left)) { return; }
274        if (bubbleUpFix(parent.right)) { return; }
275  
276        if (parent.color == DOUBLE_BLACK) 
277        { 
278           if (parent.parent == null) { parent.color = BLACK; }
279           else { bubbleUp(parent.parent); }
280        }
281     }
282  
283     /**
284        Fixes a negative-red or double-red violation introduced by bubbling up.
285        @param child the child to check for negative-red or double-red violations
286        @return true if the tree was fixed
287     */
288     private boolean bubbleUpFix(Node child)
289     {
290        if (child.color == NEGATIVE_RED) { fixNegativeRed(child); return true; }
291        else if (child.color == RED)
292        {
293           if (child.left != null && child.left.color == RED) 
294           { 
295              fixDoubleRed(child.left); return true;
296           }
297           if (child.right != null && child.right.color == RED) 
298           { 
299              fixDoubleRed(child.right); return true;
300           }
301        }
302        return false; 
303     }
304     
305     /**
306        Fixes a "double red" violation.
307        @param child the child with a red parent
308     */
309     private void fixDoubleRed(Node child)
310     {
311        Node parent = child.parent;      
312        Node grandParent = parent.parent;
313        if (grandParent == null) { parent.color = BLACK; return; }
314        Node n1, n2, n3, t1, t2, t3, t4;
315        if (parent == grandParent.left)
316        {
317           n3 = grandParent; t4 = grandParent.right;
318           if (child == parent.left)
319           {
320              n1 = child; n2 = parent;
321              t1 = child.left; t2 = child.right; t3 = parent.right;
322           }
323           else
324           {
325              n1 = parent; n2 = child;
326              t1 = parent.left; t2 = child.left; t3 = child.right; 
327           }
328        }
329        else
330        {
331           n1 = grandParent; t1 = grandParent.left;
332           if (child == parent.left)
333           {
334              n2 = child; n3 = parent;
335              t2 = child.left; t3 = child.right; t4 = parent.right;
336           }
337           else
338           {
339              n2 = parent; n3 = child;
340              t2 = parent.left; t3 = child.left; t4 = child.right; 
341           }         
342        }
343        
344        replaceWith(grandParent, n2);      
345        n1.setLeftChild(t1);
346        n1.setRightChild(t2);
347        n2.setLeftChild(n1);
348        n2.setRightChild(n3);
349        n3.setLeftChild(t3);
350        n3.setRightChild(t4);
351        n2.color = grandParent.color - 1; 
352        n1.color = BLACK;
353        n3.color = BLACK;
354  
355        if (n2 == root)
356        {
357           root.color = BLACK;
358        }
359        else if (n2.color == RED && n2.parent.color == RED)
360        {
361           fixDoubleRed(n2);
362        }
363     }
364     
365     /**
366        Fixes a "negative red" violation.
367        @param negRed the negative red node
368     */
369     private void fixNegativeRed(Node negRed)
370     {	
371        Node parent = negRed.parent;
372        Node child;
373        if (parent.left == negRed)
374        {
375           Node n1 = negRed.left;
376           Node n2 = negRed;
377           Node n3 = negRed.right;
378           Node n4 = parent;
379           Node t1 = n3.left;
380           Node t2 = n3.right;
381           Node t3 = n4.right;
382           n1.color = RED;
383           n2.color = BLACK;
384           n4.color = BLACK;
385  
386           replaceWith(n4, n3);
387           n3.setLeftChild(n2);
388           n3.setRightChild(n4);
389           n2.setLeftChild(n1);
390           n2.setRightChild(t1);
391           n4.setLeftChild(t2);
392           n4.setRightChild(t3);
393  
394           child = n1;
395        }
396        else // Mirror image
397        {
398           Node n4 = negRed.right;
399           Node n3 = negRed;
400           Node n2 = negRed.left;
401           Node n1 = parent;
402           Node t3 = n2.right;
403           Node t2 = n2.left;
404           Node t1 = n1.left;
405           n4.color = RED;
406           n3.color = BLACK;
407           n1.color = BLACK;
408  
409           replaceWith(n1, n2);
410           n2.setRightChild(n3);
411           n2.setLeftChild(n1);
412           n3.setRightChild(n4);
413           n3.setLeftChild(t3);
414           n1.setRightChild(t2);
415           n1.setLeftChild(t1);
416  
417           child = n4;
418        }
419  	   
420        if (child.left != null && child.left.color == RED) 
421        { 
422           fixDoubleRed(child.left); 
423        }
424        else if (child.right != null && child.right.color == RED) 
425        { 
426           fixDoubleRed(child.right);  
427        }
428     }
429  }
430  
431  
432