feat: add Swift codes for avl_tree article (#302)

This commit is contained in:
nuomi1 2023-01-29 00:03:28 +08:00 committed by GitHub
parent f58068c838
commit af8ee88e50
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 394 additions and 4 deletions

View file

@ -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"]),
]
)

View file

@ -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)")
}
}

View file

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