hello-algo/codes/swift/chapter_tree/avl_tree.swift
nuomi1 7359a7cb4b
Review Swift codes (#1150)
* feat(swift): review for chapter_computational_complexity

* feat(swift): review for chapter_data_structure

* feat(swift): review for chapter_array_and_linkedlist

* feat(swift): review for chapter_stack_and_queue

* feat(swift): review for chapter_hashing

* feat(swift): review for chapter_tree

* feat(swift): add codes for heap article

* feat(swift): review for chapter_heap

* feat(swift): review for chapter_graph

* feat(swift): review for chapter_searching

* feat(swift): review for chapter_sorting

* feat(swift): review for chapter_divide_and_conquer

* feat(swift): review for chapter_backtracking

* feat(swift): review for chapter_dynamic_programming

* feat(swift): review for chapter_greedy

* feat(swift): review for utils

* feat(swift): update ci tool

* feat(swift): trailing closure

* feat(swift): array init

* feat(swift): map index
2024-03-20 21:15:39 +08:00

230 lines
7.3 KiB
Swift
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

/**
* File: avl_tree.swift
* Created Time: 2023-01-28
* Author: nuomi1 (nuomi1@qq.com)
*/
import utils
/* AVL */
class AVLTree {
fileprivate var root: TreeNode? //
init() {}
/* */
func height(node: TreeNode?) -> Int {
// -1 0
node?.height ?? -1
}
/* */
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
}
/* */
func insert(val: Int) {
root = insertHelper(node: root, val: val)
}
/* */
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
}
/* */
func remove(val: Int) {
root = removeHelper(node: root, val: val)
}
/* */
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 ?? node?.right
// = 0 node
if child == nil {
return nil
}
// = 1 node
else {
node = child
}
} else {
// = 2
var temp = node?.right
while temp?.left != nil {
temp = temp?.left
}
node?.right = removeHelper(node: node?.right, val: temp!.val)
node?.val = temp!.val
}
}
updateHeight(node: node) //
/* 2. 使 */
node = rotate(node: node)
//
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)")
}
}