diff --git a/codes/swift/Package.swift b/codes/swift/Package.swift index 1dc640324..ad0354755 100644 --- a/codes/swift/Package.swift +++ b/codes/swift/Package.swift @@ -26,6 +26,7 @@ let package = Package( .executable(name: "binary_tree_bfs", targets: ["binary_tree_bfs"]), .executable(name: "binary_tree_dfs", targets: ["binary_tree_dfs"]), .executable(name: "binary_search_tree", targets: ["binary_search_tree"]), + .executable(name: "avl_tree", targets: ["avl_tree"]), ], targets: [ .target(name: "utils", path: "utils"), @@ -50,5 +51,6 @@ let package = Package( .executableTarget(name: "binary_tree_bfs", dependencies: ["utils"], path: "chapter_tree", sources: ["binary_tree_bfs.swift"]), .executableTarget(name: "binary_tree_dfs", dependencies: ["utils"], path: "chapter_tree", sources: ["binary_tree_dfs.swift"]), .executableTarget(name: "binary_search_tree", dependencies: ["utils"], path: "chapter_tree", sources: ["binary_search_tree.swift"]), + .executableTarget(name: "avl_tree", dependencies: ["utils"], path: "chapter_tree", sources: ["avl_tree.swift"]), ] ) diff --git a/codes/swift/chapter_tree/avl_tree.swift b/codes/swift/chapter_tree/avl_tree.swift new file mode 100644 index 000000000..f0ea2fcaa --- /dev/null +++ b/codes/swift/chapter_tree/avl_tree.swift @@ -0,0 +1,242 @@ +/** + * File: avl_tree.swift + * Created Time: 2023-01-28 + * Author: nuomi1 (nuomi1@qq.com) + */ + +import utils + +// Tree class +class AVLTree { + fileprivate var root: TreeNode? // 根节点 + + /* 获取结点高度 */ + func height(node: TreeNode?) -> Int { + // 空结点高度为 -1 ,叶结点高度为 0 + node == nil ? -1 : node!.height + } + + /* 更新结点高度 */ + private func updateHeight(node: TreeNode?) { + // 结点高度等于最高子树高度 + 1 + node?.height = max(height(node: node?.left), height(node: node?.right)) + 1 + } + + /* 获取平衡因子 */ + func balanceFactor(node: TreeNode?) -> Int { + // 空结点平衡因子为 0 + guard let node = node else { return 0 } + // 结点平衡因子 = 左子树高度 - 右子树高度 + return height(node: node.left) - height(node: node.right) + } + + /* 右旋操作 */ + private func rightRotate(node: TreeNode?) -> TreeNode? { + let child = node?.left + let grandChild = child?.right + // 以 child 为原点,将 node 向右旋转 + child?.right = node + node?.left = grandChild + // 更新结点高度 + updateHeight(node: node) + updateHeight(node: child) + // 返回旋转后子树的根节点 + return child + } + + /* 左旋操作 */ + private func leftRotate(node: TreeNode?) -> TreeNode? { + let child = node?.right + let grandChild = child?.left + // 以 child 为原点,将 node 向左旋转 + child?.left = node + node?.right = grandChild + // 更新结点高度 + updateHeight(node: node) + updateHeight(node: child) + // 返回旋转后子树的根节点 + return child + } + + /* 执行旋转操作,使该子树重新恢复平衡 */ + private func rotate(node: TreeNode?) -> TreeNode? { + // 获取结点 node 的平衡因子 + let balanceFactor = balanceFactor(node: node) + // 左偏树 + if balanceFactor > 1 { + if self.balanceFactor(node: node?.left) >= 0 { + // 右旋 + return rightRotate(node: node) + } else { + // 先左旋后右旋 + node?.left = leftRotate(node: node?.left) + return rightRotate(node: node) + } + } + // 右偏树 + if balanceFactor < -1 { + if self.balanceFactor(node: node?.right) <= 0 { + // 左旋 + return leftRotate(node: node) + } else { + // 先右旋后左旋 + node?.right = rightRotate(node: node?.right) + return leftRotate(node: node) + } + } + // 平衡树,无需旋转,直接返回 + return node + } + + /* 插入结点 */ + @discardableResult + func insert(val: Int) -> TreeNode? { + root = insertHelper(node: root, val: val) + return root + } + + /* 递归插入结点(辅助函数) */ + private func insertHelper(node: TreeNode?, val: Int) -> TreeNode? { + var node = node + if node == nil { + return TreeNode(x: val) + } + /* 1. 查找插入位置,并插入结点 */ + if val < node!.val { + node?.left = insertHelper(node: node?.left, val: val) + } else if val > node!.val { + node?.right = insertHelper(node: node?.right, val: val) + } else { + return node // 重复结点不插入,直接返回 + } + updateHeight(node: node) // 更新结点高度 + /* 2. 执行旋转操作,使该子树重新恢复平衡 */ + node = rotate(node: node) + // 返回子树的根节点 + return node + } + + /* 删除结点 */ + @discardableResult + func remove(val: Int) -> TreeNode? { + root = removeHelper(node: root, val: val) + return root + } + + /* 递归删除结点(辅助函数) */ + private func removeHelper(node: TreeNode?, val: Int) -> TreeNode? { + var node = node + if node == nil { + return nil + } + /* 1. 查找结点,并删除之 */ + if val < node!.val { + node?.left = removeHelper(node: node?.left, val: val) + } else if val > node!.val { + node?.right = removeHelper(node: node?.right, val: val) + } else { + if node?.left == nil || node?.right == nil { + let child = node?.left != nil ? node?.left : node?.right + // 子结点数量 = 0 ,直接删除 node 并返回 + if child == nil { + return nil + } + // 子结点数量 = 1 ,直接删除 node + else { + node = child + } + } else { + // 子结点数量 = 2 ,则将中序遍历的下个结点删除,并用该结点替换当前结点 + let temp = getInOrderNext(node: node?.right) + node?.right = removeHelper(node: node?.right, val: temp!.val) + node?.val = temp!.val + } + } + updateHeight(node: node) // 更新结点高度 + /* 2. 执行旋转操作,使该子树重新恢复平衡 */ + node = rotate(node: node) + // 返回子树的根节点 + return node + } + + /* 获取中序遍历中的下一个结点(仅适用于 root 有左子结点的情况) */ + private func getInOrderNext(node: TreeNode?) -> TreeNode? { + var node = node + if node == nil { + return node + } + // 循环访问左子结点,直到叶结点时为最小结点,跳出 + while node?.left != nil { + node = node?.left + } + return node + } + + /* 查找结点 */ + func search(val: Int) -> TreeNode? { + var cur = root + while cur != nil { + // 目标结点在 cur 的右子树中 + if cur!.val < val { + cur = cur?.right + } + // 目标结点在 cur 的左子树中 + else if cur!.val > val { + cur = cur?.left + } + // 找到目标结点,跳出循环 + else { + break + } + } + // 返回目标结点 + return cur + } +} + +@main +enum _AVLTree { + static func testInsert(tree: AVLTree, val: Int) { + tree.insert(val: val) + print("\n插入结点 \(val) 后,AVL 树为") + PrintUtil.printTree(root: tree.root) + } + + static func testRemove(tree: AVLTree, val: Int) { + tree.remove(val: val) + print("\n删除结点 \(val) 后,AVL 树为") + PrintUtil.printTree(root: tree.root) + } + + /* Driver Code */ + static func main() { + /* 初始化空 AVL 树 */ + let avlTree = AVLTree() + + /* 插入结点 */ + // 请关注插入结点后,AVL 树是如何保持平衡的 + testInsert(tree: avlTree, val: 1) + testInsert(tree: avlTree, val: 2) + testInsert(tree: avlTree, val: 3) + testInsert(tree: avlTree, val: 4) + testInsert(tree: avlTree, val: 5) + testInsert(tree: avlTree, val: 8) + testInsert(tree: avlTree, val: 7) + testInsert(tree: avlTree, val: 9) + testInsert(tree: avlTree, val: 10) + testInsert(tree: avlTree, val: 6) + + /* 插入重复结点 */ + testInsert(tree: avlTree, val: 7) + + /* 删除结点 */ + // 请关注删除结点后,AVL 树是如何保持平衡的 + testRemove(tree: avlTree, val: 8) // 删除度为 0 的结点 + testRemove(tree: avlTree, val: 5) // 删除度为 1 的结点 + testRemove(tree: avlTree, val: 4) // 删除度为 2 的结点 + + /* 查询结点 */ + let node = avlTree.search(val: 7) + print("\n查找到的结点对象为 \(node!),结点值 = \(node!.val)") + } +} diff --git a/docs/chapter_tree/avl_tree.md b/docs/chapter_tree/avl_tree.md index 6f6272645..95c818c0c 100644 --- a/docs/chapter_tree/avl_tree.md +++ b/docs/chapter_tree/avl_tree.md @@ -103,7 +103,18 @@ G. M. Adelson-Velsky 和 E. M. Landis 在其 1962 年发表的论文 "An algorit === "Swift" ```swift title="avl_tree.swift" + /* AVL 树结点类 */ + class TreeNode { + var val: Int // 结点值 + var height: Int // 结点高度 + var left: TreeNode? // 左子结点 + var right: TreeNode? // 右子结点 + init(x: Int) { + val = x + height = 0 + } + } ``` 「结点高度」是最远叶结点到该结点的距离,即走过的「边」的数量。需要特别注意,**叶结点的高度为 0 ,空结点的高度为 -1**。我们封装两个工具函数,分别用于获取与更新结点的高度。 @@ -210,7 +221,17 @@ G. M. Adelson-Velsky 和 E. M. Landis 在其 1962 年发表的论文 "An algorit === "Swift" ```swift title="avl_tree.swift" + /* 获取结点高度 */ + func height(node: TreeNode?) -> Int { + // 空结点高度为 -1 ,叶结点高度为 0 + node == nil ? -1 : node!.height + } + /* 更新结点高度 */ + func updateHeight(node: TreeNode?) { + // 结点高度等于最高子树高度 + 1 + node?.height = max(height(node: node?.left), height(node: node?.right)) + 1 + } ``` ### 结点平衡因子 @@ -295,7 +316,13 @@ G. M. Adelson-Velsky 和 E. M. Landis 在其 1962 年发表的论文 "An algorit === "Swift" ```swift title="avl_tree.swift" - + /* 获取平衡因子 */ + func balanceFactor(node: TreeNode?) -> Int { + // 空结点平衡因子为 0 + guard let node = node else { return 0 } + // 结点平衡因子 = 左子树高度 - 右子树高度 + return height(node: node.left) - height(node: node.right) + } ``` !!! note @@ -427,7 +454,19 @@ AVL 树的独特之处在于「旋转 Rotation」的操作,其可 **在不影 === "Swift" ```swift title="avl_tree.swift" - + /* 右旋操作 */ + func rightRotate(node: TreeNode?) -> TreeNode? { + let child = node?.left + let grandChild = child?.right + // 以 child 为原点,将 node 向右旋转 + child?.right = node + node?.left = grandChild + // 更新结点高度 + updateHeight(node: node) + updateHeight(node: child) + // 返回旋转后子树的根节点 + return child + } ``` ### Case 2 - 左旋 @@ -541,7 +580,19 @@ AVL 树的独特之处在于「旋转 Rotation」的操作,其可 **在不影 === "Swift" ```swift title="avl_tree.swift" - + /* 左旋操作 */ + func leftRotate(node: TreeNode?) -> TreeNode? { + let child = node?.right + let grandChild = child?.left + // 以 child 为原点,将 node 向左旋转 + child?.left = node + node?.right = grandChild + // 更新结点高度 + updateHeight(node: node) + updateHeight(node: child) + // 返回旋转后子树的根节点 + return child + } ``` ### Case 3 - 先左后右 @@ -745,7 +796,35 @@ AVL 树的独特之处在于「旋转 Rotation」的操作,其可 **在不影 === "Swift" ```swift title="avl_tree.swift" - + /* 执行旋转操作,使该子树重新恢复平衡 */ + func rotate(node: TreeNode?) -> TreeNode? { + // 获取结点 node 的平衡因子 + let balanceFactor = balanceFactor(node: node) + // 左偏树 + if balanceFactor > 1 { + if self.balanceFactor(node: node?.left) >= 0 { + // 右旋 + return rightRotate(node: node) + } else { + // 先左旋后右旋 + node?.left = leftRotate(node: node?.left) + return rightRotate(node: node) + } + } + // 右偏树 + if balanceFactor < -1 { + if self.balanceFactor(node: node?.right) <= 0 { + // 左旋 + return leftRotate(node: node) + } else { + // 先右旋后左旋 + node?.right = rightRotate(node: node?.right) + return leftRotate(node: node) + } + } + // 平衡树,无需旋转,直接返回 + return node + } ``` ## AVL 树常用操作 @@ -894,7 +973,33 @@ AVL 树的独特之处在于「旋转 Rotation」的操作,其可 **在不影 === "Swift" ```swift title="avl_tree.swift" + /* 插入结点 */ + @discardableResult + func insert(val: Int) -> TreeNode? { + root = insertHelper(node: root, val: val) + return root + } + /* 递归插入结点(辅助函数) */ + func insertHelper(node: TreeNode?, val: Int) -> TreeNode? { + var node = node + if node == nil { + return TreeNode(x: val) + } + /* 1. 查找插入位置,并插入结点 */ + if val < node!.val { + node?.left = insertHelper(node: node?.left, val: val) + } else if val > node!.val { + node?.right = insertHelper(node: node?.right, val: val) + } else { + return node // 重复结点不插入,直接返回 + } + updateHeight(node: node) // 更新结点高度 + /* 2. 执行旋转操作,使该子树重新恢复平衡 */ + node = rotate(node: node) + // 返回子树的根节点 + return node + } ``` ### 删除结点 @@ -1100,7 +1205,48 @@ AVL 树的独特之处在于「旋转 Rotation」的操作,其可 **在不影 === "Swift" ```swift title="avl_tree.swift" + /* 删除结点 */ + @discardableResult + func remove(val: Int) -> TreeNode? { + root = removeHelper(node: root, val: val) + return root + } + /* 递归删除结点(辅助函数) */ + func removeHelper(node: TreeNode?, val: Int) -> TreeNode? { + var node = node + if node == nil { + return nil + } + /* 1. 查找结点,并删除之 */ + if val < node!.val { + node?.left = removeHelper(node: node?.left, val: val) + } else if val > node!.val { + node?.right = removeHelper(node: node?.right, val: val) + } else { + if node?.left == nil || node?.right == nil { + let child = node?.left != nil ? node?.left : node?.right + // 子结点数量 = 0 ,直接删除 node 并返回 + if child == nil { + return nil + } + // 子结点数量 = 1 ,直接删除 node + else { + node = child + } + } else { + // 子结点数量 = 2 ,则将中序遍历的下个结点删除,并用该结点替换当前结点 + let temp = getInOrderNext(node: node?.right) + node?.right = removeHelper(node: node?.right, val: temp!.val) + node?.val = temp!.val + } + } + updateHeight(node: node) // 更新结点高度 + /* 2. 执行旋转操作,使该子树重新恢复平衡 */ + node = rotate(node: node) + // 返回子树的根节点 + return node + } ``` ### 查找结点