Improve removing node with two children

Previously, if the node had two children we'd remove one of the children and copy its value into the node you called remove() on. The advantage of that approach is that calling remove() on the root of the tree doesn't invalidate the root node. However, this wasn't true for when you tried to delete a root node that has only one child. That was confusing.

Now we always truly remove the node you called remove() on, and tell you what its replacement is. As a consequence, if you remove the root node you must make sure to use its replacement node as the new root from then on.

Author: hollance

Binary Search Tree/README.markdown

@@ -431,34 +431,62 @@ It returns the rightmost descendent of the node. We find it by following `right`
 Finally, we can write the code the remove a node from the tree:
 
 ```swift
-  public func remove() {
+  public func remove() -> BinarySearchTree? {
+    let replacement: BinarySearchTree?
+
     if let left = left {
       if let right = right {
-        let successor = right.minimum()  // 1
-        value = successor.value
-        successor.remove()
+        replacement = removeNodeWithTwoChildren(left, right)  // 1
       } else {
-        reconnectParentToNode(left)      // 2
+        replacement = left           // 2
       }
-    } else if let right = right {        // 3
-      reconnectParentToNode(right)
+    } else if let right = right {    // 3
+      replacement = right
     } else {
-      reconnectParentToNode(nil)         // 4 
+      replacement = nil              // 4
     }
+    
+    reconnectParentToNode(replacement)
+
+    parent = nil
+    left = nil
+    right = nil
+
+    return replacement
   }
 ```
 
-It doesn't look so scary after all. ;-) Here is what it does:
-
-1. This node has two children. It must be replaced by the smallest child that is larger than this node's value, which is the leftmost descendent of the right child, i.e. `right.minimum()`. 
+It doesn't look so scary after all. ;-) There are four situations to handle:
 
+1. This node has two children. 
 2. This node only has a left child. The left child replaces the node.
-
 3. This node only has a right child. The right child replaces the node.
-
 4. This node has no children. We just disconnect it from its parent.
 
-The only tricky situation here is `// 1`. Rather than deleting the current node, we give it the successor's value and remove the successor node instead.
+First, we determine which node will replace the one we're removing and then we call `reconnectParentToNode()` to change the `left`, `right`, and `parent` pointers to make that happen. Since the current node is no longer part of the tree, we clean it up by setting its pointers to `nil`. Finally, we return the node that has replaced the removed one (or `nil` if this was a leaf node).
+
+The only tricky situation here is `// 1` and that logic has its own helper method:
+
+```swift
+  private func removeNodeWithTwoChildren(left: BinarySearchTree, _ right: BinarySearchTree) -> BinarySearchTree {
+    let successor = right.minimum()
+    successor.remove()
+    
+    successor.left = left
+    left.parent = successor
+    
+    if right !== successor {
+      successor.right = right
+      right.parent = successor
+    } else {
+      successor.right = nil
+    }
+    
+    return successor
+  }
+```
+
+If the node to remove has two children, it must be replaced by the smallest child that is larger than this node's value. That always happens to be the leftmost descendent of the right child, i.e. `right.minimum()`. We take that node out of its original position in the tree and put it into the place of the node we're removing.
 
 Try it out:
 
@@ -478,9 +506,9 @@ But after `remove()` you get:
 
 	((1) <- 5) <- 7 -> ((9) <- 10)
 
-As you can see, node `5` has taken the place of `2`. In fact, if you do `print(node2)` you'll see that it now has the value `5`. We didn't actually remove the `node2` object, we just gave it a new value.
+As you can see, node `5` has taken the place of `2`.
 
-> **Note:** What would happen if you deleted the root node? Specifically, how would you know which node becomes the new root node? It turns out you don't have to worry about that because the root never actually gets deleted. We simply give it a new value.
+> **Note:** What would happen if you deleted the root node? In that case, `remove()` tells you which node has become the new root. Try it out: call `tree.remove()` and see what happens.
 
 Like most binary search tree operations, removing a node runs in **O(h)** time, where **h** is the height of the tree.
 

Binary Search Tree/Solution 1/BinarySearchTree.playground/Contents.swift

@@ -26,7 +26,7 @@ tree.maximum()
 
 if let node2 = tree.search(2) {
   node2.remove()
-  node2.value          // this is now node "5"
+  node2
   print(tree)
 }
 

Binary Search Tree/Solution 1/BinarySearchTree.playground/Sources/BinarySearchTree.swift

@@ -101,34 +101,67 @@ extension BinarySearchTree {
 extension BinarySearchTree {
   /*
     Deletes a node from the tree.
+
+    Returns the node that has replaced this removed one (or nil if this was a
+    leaf node). That is primarily useful for when you delete the root node, in
+    which case the tree gets a new root.
+
     Performance: runs in O(h) time, where h is the height of the tree.
   */
-  public func remove() {
+  public func remove() -> BinarySearchTree? {
+    let replacement: BinarySearchTree?
+
     if let left = left {
       if let right = right {
-        // This node has two children. It must be replaced by the smallest
-        // child that is larger than this node's value, which is the leftmost
-        // descendent of the right child.
-        let successor = right.minimum()
-        
-        // Rather than deleting the current node (which is problematic for the
-        // root node) we give it the successor's value and remove the successor.
-        value = successor.value
-        
-        // If this in-order successor has a right child of its own (it cannot
-        // have a left child by definition), then that must take its place.
-        successor.remove()
+        replacement = removeNodeWithTwoChildren(left, right)
       } else {
         // This node only has a left child. The left child replaces the node.
-        reconnectParentToNode(left)
+        replacement = left
       }
     } else if let right = right {
       // This node only has a right child. The right child replaces the node.
-      reconnectParentToNode(right)
+      replacement = right
     } else {
       // This node has no children. We just disconnect it from its parent.
-      reconnectParentToNode(nil)
+      replacement = nil
+    }
+
+    reconnectParentToNode(replacement)
+
+    // The current node is no longer part of the tree, so clean it up.
+    parent = nil
+    left = nil
+    right = nil
+
+    return replacement
+  }
+
+  private func removeNodeWithTwoChildren(left: BinarySearchTree, _ right: BinarySearchTree) -> BinarySearchTree {
+    // This node has two children. It must be replaced by the smallest
+    // child that is larger than this node's value, which is the leftmost
+    // descendent of the right child.
+    let successor = right.minimum()
+
+    // If this in-order successor has a right child of its own (it cannot
+    // have a left child by definition), then that must take its place.
+    successor.remove()
+
+    // Connect our left child with the new node.
+    successor.left = left
+    left.parent = successor
+
+    // Connect our right child with the new node. If the right child does
+    // not have any left children of its own, then the in-order successor
+    // *is* the right child.
+    if right !== successor {
+      successor.right = right
+      right.parent = successor
+    } else {
+      successor.right = nil
     }
+
+    // And finally, connect the successor node to our parent.
+    return successor
   }
 
   private func reconnectParentToNode(node: BinarySearchTree?) {

Binary Search Tree/Solution 1/BinarySearchTree.swift

@@ -101,34 +101,67 @@ extension BinarySearchTree {
 extension BinarySearchTree {
   /*
     Deletes a node from the tree.
+
+    Returns the node that has replaced this removed one (or nil if this was a
+    leaf node). That is primarily useful for when you delete the root node, in
+    which case the tree gets a new root.
+
     Performance: runs in O(h) time, where h is the height of the tree.
   */
-  public func remove() {
+  public func remove() -> BinarySearchTree? {
+    let replacement: BinarySearchTree?
+
     if let left = left {
       if let right = right {
-        // This node has two children. It must be replaced by the smallest
-        // child that is larger than this node's value, which is the leftmost
-        // descendent of the right child.
-        let successor = right.minimum()
-        
-        // Rather than deleting the current node (which is problematic for the
-        // root node) we give it the successor's value and remove the successor.
-        value = successor.value
-        
-        // If this in-order successor has a right child of its own (it cannot
-        // have a left child by definition), then that must take its place.
-        successor.remove()
+        replacement = removeNodeWithTwoChildren(left, right)
       } else {
         // This node only has a left child. The left child replaces the node.
-        reconnectParentToNode(left)
+        replacement = left
       }
     } else if let right = right {
       // This node only has a right child. The right child replaces the node.
-      reconnectParentToNode(right)
+      replacement = right
     } else {
       // This node has no children. We just disconnect it from its parent.
-      reconnectParentToNode(nil)
+      replacement = nil
+    }
+
+    reconnectParentToNode(replacement)
+
+    // The current node is no longer part of the tree, so clean it up.
+    parent = nil
+    left = nil
+    right = nil
+
+    return replacement
+  }
+
+  private func removeNodeWithTwoChildren(left: BinarySearchTree, _ right: BinarySearchTree) -> BinarySearchTree {
+    // This node has two children. It must be replaced by the smallest
+    // child that is larger than this node's value, which is the leftmost
+    // descendent of the right child.
+    let successor = right.minimum()
+
+    // If this in-order successor has a right child of its own (it cannot
+    // have a left child by definition), then that must take its place.
+    successor.remove()
+
+    // Connect our left child with the new node.
+    successor.left = left
+    left.parent = successor
+
+    // Connect our right child with the new node. If the right child does
+    // not have any left children of its own, then the in-order successor
+    // *is* the right child.
+    if right !== successor {
+      successor.right = right
+      right.parent = successor
+    } else {
+      successor.right = nil
     }
+
+    // And finally, connect the successor node to our parent.
+    return successor
   }
 
   private func reconnectParentToNode(node: BinarySearchTree?) {

Binary Search Tree/Solution 1/Tests/BinarySearchTreeTests.swift

@@ -139,14 +139,17 @@ class BinarySearchTreeTest: XCTestCase {
     XCTAssertTrue(node5.left === node4)
     XCTAssertNil(node5.right)
 
-    node4.remove()
+    let replacement1 = node4.remove()
     XCTAssertNil(node5.left)
+    XCTAssertNil(replacement1)
 
-    node5.remove()
+    let replacement2 = node5.remove()
     XCTAssertNil(tree.left)
+    XCTAssertNil(replacement2)
 
-    node10.remove()
+    let replacement3 = node10.remove()
     XCTAssertNil(tree.right)
+    XCTAssertNil(replacement3)
 
     XCTAssertEqual(tree.count, 1)
     XCTAssertEqual(tree.toArray(), [8])
@@ -223,15 +226,18 @@ class BinarySearchTreeTest: XCTestCase {
     XCTAssertTrue(node5 === node4.parent)
     XCTAssertTrue(node5 === node6.parent)
 
-    node5.remove()
-    XCTAssertTrue(tree.left === node5)
-    XCTAssertTrue(tree === node5.parent)
-    XCTAssertTrue(node5.left === node4)
-    XCTAssertTrue(node5 === node4.parent)
+    let replacement1 = node5.remove()
+    XCTAssertTrue(replacement1 === node6)
+    XCTAssertTrue(tree.left === node6)
+    XCTAssertTrue(tree === node6.parent)
+    XCTAssertTrue(node6.left === node4)
+    XCTAssertTrue(node6 === node4.parent)
+    XCTAssertNil(node5.left)
     XCTAssertNil(node5.right)
+    XCTAssertNil(node5.parent)
     XCTAssertNil(node4.left)
     XCTAssertNil(node4.right)
-    XCTAssertEqual(node5.value, 6)
+    XCTAssertNotNil(node4.parent)
     XCTAssertEqual(tree.count, 6)
     XCTAssertEqual(tree.toArray(), [4, 6, 8, 9, 10, 11])
 
@@ -243,15 +249,18 @@ class BinarySearchTreeTest: XCTestCase {
     XCTAssertTrue(node10 === node9.parent)
     XCTAssertTrue(node10 === node11.parent)
 
-    node10.remove()
-    XCTAssertTrue(tree.right === node10)
-    XCTAssertTrue(tree === node10.parent)
-    XCTAssertTrue(node10.left === node9)
-    XCTAssertTrue(node10 === node9.parent)
+    let replacement2 = node10.remove()
+    XCTAssertTrue(replacement2 === node11)
+    XCTAssertTrue(tree.right === node11)
+    XCTAssertTrue(tree === node11.parent)
+    XCTAssertTrue(node11.left === node9)
+    XCTAssertTrue(node11 === node9.parent)
+    XCTAssertNil(node10.left)
     XCTAssertNil(node10.right)
+    XCTAssertNil(node10.parent)
     XCTAssertNil(node9.left)
     XCTAssertNil(node9.right)
-    XCTAssertEqual(node10.value, 11)
+    XCTAssertNotNil(node9.parent)
     XCTAssertEqual(tree.count, 5)
     XCTAssertEqual(tree.toArray(), [4, 6, 8, 9, 11])
   }
@@ -273,28 +282,39 @@ class BinarySearchTreeTest: XCTestCase {
     XCTAssertTrue(node11.right === node15)
     XCTAssertTrue(node11 === node15.parent)
 
-    node10.remove()
-    XCTAssertTrue(tree.right === node10)
-    XCTAssertTrue(tree === node10.parent)
-    XCTAssertTrue(node10.left === node9)
-    XCTAssertTrue(node10 === node9.parent)
-    XCTAssertTrue(node10.right === node20)
-    XCTAssertTrue(node10 === node20.parent)
+    let replacement = node10.remove()
+    XCTAssertTrue(replacement === node11)
+    XCTAssertTrue(tree.right === node11)
+    XCTAssertTrue(tree === node11.parent)
+    XCTAssertTrue(node11.left === node9)
+    XCTAssertTrue(node11 === node9.parent)
+    XCTAssertTrue(node11.right === node20)
+    XCTAssertTrue(node11 === node20.parent)
     XCTAssertTrue(node20.left === node15)
     XCTAssertTrue(node20 === node15.parent)
     XCTAssertNil(node20.right)
-    XCTAssertEqual(node10.value, 11)
+    XCTAssertNil(node10.left)
+    XCTAssertNil(node10.right)
+    XCTAssertNil(node10.parent)
     XCTAssertEqual(tree.count, 8)
     XCTAssertEqual(tree.toArray(), [4, 5, 8, 9, 11, 13, 15, 20])
   }
 
   func testRemoveRoot() {
     let tree = BinarySearchTree(array: [8, 5, 10, 4, 9, 20, 11, 15, 13])
-    tree.remove()
-
-    XCTAssertEqual(tree.value, 9)
-    XCTAssertEqual(tree.count, 8)
-    XCTAssertEqual(tree.toArray(), [4, 5, 9, 10, 11, 13, 15, 20])
+    
+    let node9 = tree.search(9)!
+    
+    let newRoot = tree.remove()
+    XCTAssertTrue(newRoot === node9)
+    XCTAssertEqual(newRoot!.value, 9)
+    XCTAssertEqual(newRoot!.count, 8)
+    XCTAssertEqual(newRoot!.toArray(), [4, 5, 9, 10, 11, 13, 15, 20])
+
+    // The old root is a subtree of a single element.
+    XCTAssertEqual(tree.value, 8)
+    XCTAssertEqual(tree.count, 1)
+    XCTAssertEqual(tree.toArray(), [8])
   }
 
   func testPredecessor() {