Swift 5.3 —— 堆数据结构 Heap

堆数据结构

堆形数据结构是一个二叉树，可以通过数组构造。
堆分为最大堆和最小堆：
最大堆节点的值比子节点的值更大，根节点的值最大，
最小堆节点的值比子节点的值更小，根节点的值最小。

堆结构有很多用途，可以用来进行堆排序，可以方便计算最大值或者最小值，可以构建优先级队列，还可以构造图算法。

struct Heap<Element: Equatable> {
    var elements: [Element] = []
    let sort: (Element, Element) -> Bool
    init(sort: @escaping (Element, Element) -> Bool, elements: [Element]) {
        self.sort = sort
        self.elements = elements
        
        if !elements.isEmpty {
            for i in stride(from: elements.count / 2 - 1, through: 0, by: -1) {
                siftDown(from: i)
            }
        }
    }
    
    var isEmpty: Bool {
        elements.isEmpty
    }
    
    var count: Int {
        elements.count
    }
    
    func peek() -> Element? {
        elements.first
    }
    
    func leftChildIndex(ofParentAt index: Int) -> Int {
        (2 * index) + 1
    }
    
    func rightChildIndex(ofParentAt index: Int) -> Int {
        (2 * index) + 2
    }
    
    func parentIndex(ofChildAt index: Int) -> Int {
        (index - 1) / 2
    }
}

用数组表示一个堆型结构，需要提供顺序计算方式sort和所有的数据elements
在堆初始化的时候，就要对数组进行排序，使其符合堆型结构，即最大堆或者最小堆。

extension Heap {
    mutating func remove() -> Element? {
        guard !isEmpty else { return nil }
        elements.swapAt(0, count - 1)
        defer {
            siftDown(from: 0)
        }
        return elements.removeLast()
    }
    
    mutating func siftDown(from index: Int) {
        var parent = index
        while true {
            let left = leftChildIndex(ofParentAt: parent)
            let right = rightChildIndex(ofParentAt: parent)
            var candidate = parent
            if left < count, sort(elements[left], elements[candidate]) {
                candidate = left
            }
            if right < count, sort(elements[right], elements[candidate]) {
                candidate = right
            }
            if candidate == parent {
                return
            }
            elements.swapAt(parent, candidate)
            parent = candidate
        }
    }
    
    mutating func insert(_ element: Element) {
        elements.append(element)
        siftUp(from: elements.count - 1)
    }
    
    mutating func siftUp(from index: Int) {
        var child = index
        var parent = parentIndex(ofChildAt: child)
        while child > 0, sort(elements[child], elements[parent]) {
            elements.swapAt(child, parent)
            parent = parentIndex(ofChildAt: child)
        }
    }
    
    mutating func remove(at index: Int) -> Element? {
        guard index < elements.count else {
            return nil
        }
        if index == elements.count - 1 {
            return elements.removeLast()
        } else {
            elements.swapAt(index, elements.count - 1)
            defer {
                siftDown(from: index)
                siftDown(from: index)
            }
            return elements.removeLast()
        }
    }
    
    func index(of element: Element, startingAt i: Int) -> Int? {
        if i >= count {
            return nil
        }
        if sort(element, elements[i]) {
            return nil
        }
        if element == elements[i] {
            return i
        }
        if let j = index(of: element, startingAt: leftChildIndex(ofParentAt: i)) {
            return j
        }
        
        if let j = index(of: element, startingAt: rightChildIndex(ofParentAt: i)) {
            return j
        }
        return nil
    }
}

一个完整的堆型结构，还需要包含插入，删除，查询等常用方法。

github仓库地址 https://github.com/SunZhiC/DataStructuresInSwift

References

Data Structures & Algorithms in Swift by Matthijs Hollemans.
https://github.com/apple/swift-algorithms
https://github.com/raywenderlich/swift-algorithm-club