该板块只是为了记录一些学习时的重点，如有看官想看深入的讲解，请移步swift算法俱乐部
AVLNode.swift
public class AVLNode<Element> {
public var value: Element
public var leftChild: AVLNode?
public var rightChild: AVLNode?
public var height = 0
public var balanceFactor: Int {
return leftHeight - rightHeight
}
public var leftHeight: Int {
return leftChild?.height ?? -1
}
public var rightHeight: Int {
return rightChild?.height ?? -1
}
public init(value: Element) {
self.value = value
}
}
extension AVLNode: CustomStringConvertible {
public var description: String {
return diagram(for: self)
}
private func diagram(for node: AVLNode?,
_ top: String = "",
_ root: String = "",
_ bottom: String = "") -> String {
guard let node = node else {
return root + "nil\n"
}
if node.leftChild == nil && node.rightChild == nil {
return root + "\(node.value)\n"
}
return diagram(for: node.rightChild, top + " ", top + "┌──", top + "│ ")
+ root + "\(node.value)\n"
+ diagram(for: node.leftChild, bottom + "│ ", bottom + "└──", bottom + " ")
}
}
extension AVLNode {
public func traverseInOrder(visit: (Element) -> Void) {
leftChild?.traverseInOrder(visit: visit)
visit(value)
rightChild?.traverseInOrder(visit: visit)
}
public func traversePreOrder(visit: (Element) -> Void) {
visit(value)
leftChild?.traversePreOrder(visit: visit)
rightChild?.traversePreOrder(visit: visit)
}
public func traversePostOrder(visit: (Element) -> Void) {
leftChild?.traversePostOrder(visit: visit)
rightChild?.traversePostOrder(visit: visit)
visit(value)
}
}
AVLTree.swift
public struct AVLTree<Element: Comparable> {
public private(set) var root: AVLNode<Element>?
public init() {}
}
extension AVLTree: CustomStringConvertible {
public var description: String {
guard let root = root else { return "empty tree" }
return String(describing: root)
}
}
extension AVLTree {
public mutating func insert(_ value: Element) {
root = insert(from: root, value: value)
}
// 插入一个节点
private func insert(from node: AVLNode<Element>?, value: Element) -> AVLNode<Element> {
guard let node = node else {
return AVLNode(value: value)
}
if value < node.value {
node.leftChild = insert(from: node.leftChild, value: value)
} else {
node.rightChild = insert(from: node.rightChild, value: value)
}
let balancedNode = balanced(node)
balancedNode.height = max(balancedNode.leftHeight, balancedNode.rightHeight) + 1
return balancedNode
}
// 左旋转
private func leftRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
let pivot = node.rightChild!
node.rightChild = pivot.leftChild
pivot.leftChild = node
node.height = max(node.leftHeight, node.rightHeight) + 1
pivot.height = max(pivot.leftHeight, pivot.rightHeight) + 1
return pivot
}
// 右旋转
private func rightRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
// 不应该增加可选校验嘛，虽然右旋转一定存在该左节点，但调用方使用错了呢
let pivot = node.leftChild!
node.leftChild = pivot.rightChild
pivot.rightChild = node
node.height = max(node.leftHeight, node.rightHeight) + 1
pivot.height = max(pivot.leftHeight, pivot.rightHeight) + 1
return pivot
}
// 右左旋转
private func rightLeftRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
guard let rightChild = node.rightChild else {
return node
}
node.rightChild = rightRotate(rightChild)
return leftRotate(node)
}
// 左右旋转
private func leftRightRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
guard let leftChild = node.leftChild else {
return node
}
node.leftChild = leftRotate(leftChild)
return rightRotate(node)
}
// 平衡某个节点
private func balanced(_ node: AVLNode<Element>) -> AVLNode<Element> {
switch node.balanceFactor {
case 2:
if let leftChild = node.leftChild, leftChild.balanceFactor == -1 {
return leftRightRotate(node)
} else {
return rightRotate(node)
}
case -2:
if let rightChild = node.rightChild, rightChild.balanceFactor == 1 {
return rightLeftRotate(node)
} else {
return leftRotate(node)
}
default:
return node
}
}
}
extension AVLTree {
// 是否包含某个value
public func contains(_ value: Element) -> Bool {
var current = root
while let node = current {
if node.value == value {
return true
}
if value < node.value {
current = node.leftChild
} else {
current = node.rightChild
}
}
return false
}
}
private extension AVLNode {
// 用来获得一个节点的最小子节点
var min: AVLNode {
return leftChild?.min ?? self
}
}
extension AVLTree {
// 移除一个节点
public mutating func remove(_ value: Element) {
root = remove(node: root, value: value)
}
private func remove(node: AVLNode<Element>?, value: Element) -> AVLNode<Element>? {
guard let node = node else {
return nil
}
if value == node.value {
if node.leftChild == nil && node.rightChild == nil {
return nil
}
if node.leftChild == nil {
return node.rightChild
}
if node.rightChild == nil {
return node.leftChild
}
node.value = node.rightChild!.min.value
node.rightChild = remove(node: node.rightChild, value: node.value)
} else if value < node.value {
node.leftChild = remove(node: node.leftChild, value: value)
} else {
node.rightChild = remove(node: node.rightChild, value: value)
}
let balancedNode = balanced(node)
balancedNode.height = max(balancedNode.leftHeight, balancedNode.rightHeight) + 1
return balancedNode
}
}
Key points
A self-balancing tree avoids performance degradation by performing a balancing procedure whenever you add or remove elements in the tree.
AVL trees preserve balance by readjusting parts of the tree when the tree is no longer balanced.
网友评论