From af8ee88e501e01095ce5d2a25359af6c471bf44b Mon Sep 17 00:00:00 2001
From: nuomi1 <nuomi1@qq.com>
Date: Sun, 29 Jan 2023 00:03:28 +0800
Subject: [PATCH] feat: add Swift codes for avl_tree article (#302)

---
 codes/swift/Package.swift               |   2 +
 codes/swift/chapter_tree/avl_tree.swift | 242 ++++++++++++++++++++++++
 docs/chapter_tree/avl_tree.md           | 156 ++++++++++++++-
 3 files changed, 395 insertions(+), 5 deletions(-)
 create mode 100644 codes/swift/chapter_tree/avl_tree.swift

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
+    }
     ```
 
 ### 查找结点