Fix red-black tree insertion invariants

src/data-structures/tree/BinaryTreeNode.js

@@ -100,8 +100,8 @@ export default class BinaryTreeNode {
    * @return {BinaryTreeNode}
    */
   setLeft(node) {
-    // Reset parent for left node since it is going to be detached.
-    if (this.left) {
+    // Detach only if the old child was not already reparented by a rotation.
+    if (this.left && this.left.parent === this) {
       this.left.parent = null;
     }
 
@@ -121,8 +121,8 @@ export default class BinaryTreeNode {
    * @return {BinaryTreeNode}
    */
   setRight(node) {
-    // Reset parent for right node since it is going to be detached.
-    if (this.right) {
+    // Detach only if the old child was not already reparented by a rotation.
+    if (this.right && this.right.parent === this) {
       this.right.parent = null;
     }
 

src/data-structures/tree/__test__/BinaryTreeNode.test.js

@@ -1,6 +1,22 @@
 import BinaryTreeNode from '../BinaryTreeNode';
 
 describe('BinaryTreeNode', () => {
+  it.each(['setLeft', 'setRight'])('should preserve a reparented child during %s', (setter) => {
+    const previousParent = new BinaryTreeNode(1);
+    const nextParent = new BinaryTreeNode(2);
+    const child = new BinaryTreeNode(3);
+    const replacement = new BinaryTreeNode(4);
+
+    previousParent[setter](child);
+    nextParent[setter](child);
+    previousParent[setter](replacement);
+
+    expect(child.parent).toBe(nextParent);
+    expect(replacement.parent).toBe(previousParent);
+    previousParent[setter](null);
+    expect(replacement.parent).toBeNull();
+  });
+
   it('should create node', () => {
     const node = new BinaryTreeNode();
 

src/data-structures/tree/red-black-tree/RedBlackTree.js

@@ -44,8 +44,9 @@ export default class RedBlackTree extends BinarySearchTree {
    * @param {BinarySearchTreeNode} node
    */
   balance(node) {
-    // If it is a root node then nothing to balance here.
-    if (this.nodeComparator.equal(node, this.root)) {
+    // Only a red node can violate the rule against adjacent red nodes.
+    // A black subtree root produced by a rotation needs no further balancing.
+    if (this.nodeComparator.equal(node, this.root) || this.isNodeBlack(node)) {
       return;
     }
 

src/data-structures/tree/red-black-tree/__test__/RedBlackTree.invariants.test.js

@@ -0,0 +1,65 @@
+import RedBlackTree from '../RedBlackTree';
+
+// Verify membership and every red-black invariant after each insertion, rather
+// than assuming a particular (non-unique) valid tree shape.
+const verifyTree = (tree, values) => {
+  expect(tree.root.parent).toBeNull();
+  expect(tree.isNodeBlack(tree.root)).toBe(true);
+  expect(tree.root.traverseInOrder()).toEqual([...values].sort((a, b) => a - b));
+
+  const blackHeight = (node, parent) => {
+    if (!node) {
+      return 1;
+    }
+
+    expect(node.parent).toBe(parent);
+    expect(tree.isNodeColored(node)).toBe(true);
+    if (tree.isNodeRed(node)) {
+      expect(!node.left || tree.isNodeBlack(node.left)).toBe(true);
+      expect(!node.right || tree.isNodeBlack(node.right)).toBe(true);
+    }
+
+    const leftHeight = blackHeight(node.left, node);
+    const rightHeight = blackHeight(node.right, node);
+    expect(leftHeight).toBe(rightHeight);
+    return leftHeight + Number(tree.isNodeBlack(node));
+  };
+
+  blackHeight(tree.root, null);
+};
+
+describe('RedBlackTree insertion invariants', () => {
+  it.each([
+    [[1, 3, 2, 4, 5]],
+    [[1, 4, 2, 3, 5, 6]],
+    [[6, 2, 1, 3, 7, 4, 5]],
+    [[5, 3, 4, 2, 1]],
+    [[6, 3, 5, 4, 2, 1]],
+    [[2, 6, 7, 5, 1, 4, 3]],
+  ])('should retain all nodes and colors for %j', (values) => {
+    const tree = new RedBlackTree();
+    values.forEach((value, index) => {
+      tree.insert(value);
+      verifyTree(tree, values.slice(0, index + 1));
+    });
+  });
+
+  it('should preserve the invariants for every permutation of six keys', () => {
+    const checkPermutations = (prefix, remaining) => {
+      if (!remaining.length) {
+        const tree = new RedBlackTree();
+        prefix.forEach((value, index) => {
+          tree.insert(value);
+          verifyTree(tree, prefix.slice(0, index + 1));
+        });
+        return;
+      }
+
+      remaining.forEach((value, index) => {
+        checkPermutations([...prefix, value], remaining.filter((_, i) => i !== index));
+      });
+    };
+
+    checkPermutations([], [1, 2, 3, 4, 5, 6]);
+  });
+});