使用Swift完成优先队列的实现

107 阅读2分钟

在算法里,我们经常会遇到优先队列的实现,但是Swift没有直接提供,因此在我们需要的时候,需要自己去实现它。

下边是一个实现基本思路及示例:

// 泛型结构体,处理任何类型 T
struct PriorityQueue<T> {
    // 存储元素的数组
    // 这个数组被组织成一个二叉堆的结构
    var elements = [T]()
    
    // 用于比较元素优先级的函数
    // 这是一个闭包,接受两个 T 类型的参数,返回一个 Bool 值
    // 如果第一个参数应该排在第二个参数之前,则返回 true
    let sort: (T, T) -> Bool

    // 初始化方法,接受一个比较函数
    // @escaping 表示这个闭包可能在初始化方法返回后才被调用
    init(sort: @escaping (T, T) -> Bool) {
        self.sort = sort
    }

    // 入队方法
    // mutating 关键字表示这个方法会修改结构体的状态
    mutating func enqueue(_ element: T) {
        // 将新元素添加到数组的末尾
        elements.append(element)
        // 通过向上调整来维护堆的性质
        siftUp(from: elements.count - 1)
    }

    // 出队方法
    // 返回类型是 T?,因为队列可能为空
    mutating func dequeue() -> T? {
        // 如果队列为空,直接返回 nil
        guard !elements.isEmpty else { return nil }
        // 交换第一个元素(要返回的元素)和最后一个元素
        elements.swapAt(0, elements.count - 1)
        // 移除并保存最后一个元素(原来的第一个元素)
        let element = elements.removeLast()
        // 如果队列不为空,需要向下调整第一个元素以维护堆的性质
        if !elements.isEmpty {
            siftDown(from: 0)
        }
        // 返回原来的第一个元素
        return element
    }

    // 向上调整堆的方法
    private mutating func siftUp(from index: Int) {
        var child = index
        var parent = (child - 1) / 2  // 计算父节点的索引
        // 当子节点不是根节点,且应该排在父节点之前时
        while child > 0 && sort(elements[child], elements[parent]) {
            // 交换子节点和父节点
            elements.swapAt(child, parent)
            // 更新索引,继续向上检查
            child = parent
            parent = (child - 1) / 2
        }
    }

    // 向下调整堆的方法
    private mutating func siftDown(from index: Int) {
        let count = elements.count
        var parent = index
        
        while true {
            let leftChild = 2 * parent + 1  // 左子节点的索引
            let rightChild = leftChild + 1  // 右子节点的索引
            var candidate = parent  // 假设父节点是最适合的候选者
            
            // 如果左子节点存在且应该排在候选者之前
            if leftChild < count && sort(elements[leftChild], elements[candidate]) {
                candidate = leftChild
            }
            // 如果右子节点存在且应该排在候选者之前
            if rightChild < count && sort(elements[rightChild], elements[candidate]) {
                candidate = rightChild
            }
            // 如果候选者就是父节点,说明堆的性质已经满足,退出循环
            if candidate == parent {
                return
            }
            // 交换父节点和候选者
            elements.swapAt(parent, candidate)
            // 更新父节点索引,继续向下检查
            parent = candidate
        }
    }
}