  1  import java.util.ArrayList;

  2  import java.util.List;

  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

 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;

 18     /**

 19        Constructs an empty tree.

 20     */

 21     public RedBlackTree()

 22     {

 23        root = null;

 24     }

 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     }

 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     }

 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

 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        }

 81        if (!found) { return; }

 83        // toBeRemoved contains obj

 85        // If one of the children is empty, use the other

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

 93           fixBeforeRemove(toBeRemoved);

 94           replaceWith(toBeRemoved, newChild);

 95           return;

 96        }

 98        // Neither subtree is empty

100        // Find smallest element of the right subtree

102        Node smallest = toBeRemoved.right;

103        while (smallest.left != null)

104        {

105           smallest = smallest.left;

106        }

108        // smallest contains smallest child in right subtree

110        // Move contents, unlink child

112        toBeRemoved.data = smallest.data;

113        fixBeforeRemove(smallest);

114        replaceWith(smallest, smallest.right);

115     }

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     }

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     }

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;

152        /**

153           Constructs a red node with no data.

154        */

155        public Node() {}

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        }

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        }

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     }

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     }

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     }

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

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     }

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--;

273        if (bubbleUpFix(parent.left)) { return; }

274        if (bubbleUpFix(parent.right)) { return; }

276        if (parent.color == DOUBLE_BLACK)

277        {

278           if (parent.parent == null) { parent.color = BLACK; }

279           else { bubbleUp(parent.parent); }

280        }

281     }

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     }

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        }

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;

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     }

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;

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);

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;

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);

417           child = n4;

418        }

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  }