Deletion in a binary tree

BinaryTreeDeletion.swift

import UIKit
import PlaygroundSupport

class TreeNode {
  var data: Int
  var left: TreeNode?
  var right: TreeNode?
  
  init(data: Int) {
    self.data = data
  }
}

extension TreeNode: Equatable {
  static func ==(lhs: TreeNode, rhs: TreeNode) -> Bool {
    return lhs.data == rhs.data && lhs.left == rhs.left && lhs.right == rhs.right
  }
}

class BinaryTree {
  var root: TreeNode
  
  init(root: TreeNode) {
    self.root = root
  }
  
  func postOrderTraversal() -> [TreeNode] {
    var result = [TreeNode]()
    if let left = root.left {
      result.append(contentsOf: BinaryTree(root: left).postOrderTraversal())
    }
    if let right = root.right {
      result.append(contentsOf: BinaryTree(root: right).postOrderTraversal())
    }
    result.append(root)
    return result
  }
    
  /// breadth first traversal
  func levelOrderTraversal() -> [Int] {
    var result = [Int]()
    var queue = [TreeNode?]()
    queue.append(root)
    while !queue.isEmpty {
      guard let node = queue.removeFirst() else { continue }
      result.append(node.data)
      queue.append(node.left)
      queue.append(node.right)
    }
    return result
  }
  
  ///delete a node from it by making sure that tree shrinks from the bottom (i.e. the deleted node is replaced by bottom most and rightmost node)
  func delete(node targetNode: TreeNode) {
    var lastNode: TreeNode?
    
    ///do a level order traversal to find deepest element since the deepest element will be the last element
    var queue = [TreeNode?]()
    queue.append(root)
    while !queue.isEmpty {
      guard let node = queue.removeFirst() else { continue }
      lastNode = node
      queue.append(node.left)
      queue.append(node.right)
    }
    
    ///find parent of the lastNode by doing a post order traversal since element next to the node found above will be its parent
    var lastNodeParent: TreeNode?
    let nodes = postOrderTraversal()
    for (index, node) in nodes.enumerated() {
      if node == lastNode {
        lastNodeParent = nodes[index + 1]
        break
      }
    }
    
    ///set parent child to null
    if lastNodeParent?.left == lastNode {
      lastNodeParent?.left = nil
    } else if lastNodeParent?.right == lastNode {
      lastNodeParent?.right = nil
    }
    
    ///copy data of the last node in the target node which needs to be deleted, this will delete the last node, effectively deleting the target node
    targetNode.data = lastNode?.data ?? -1
  }
}

let node10 = TreeNode(data: 10)
let node11 = TreeNode(data: 11)
let node9 = TreeNode(data: 9)
let node7 = TreeNode(data: 7)
let node15 = TreeNode(data: 15)
let node8 = TreeNode(data: 8)

node10.left = node11
node10.right = node9
node11.left = node7
node9.left = node15
node9.right = node8

let tree = BinaryTree(root: node10)
print("levelOrder: ", tree.levelOrderTraversal())

tree.delete(node: node9)
print("levelOrder after deletion: ", tree.levelOrderTraversal())