Skip to content
Ivan edited this page Jan 14, 2022 · 5 revisions
Clone this wiki locally

模版

优先队列

class Pq<E> {
    class PqN<E>: Comparable {
        let p: Int
        let v: E
        init(_ p: Int, _ v: E) {
            self.p = p
            self.v = v
        }

        public static func == (_ l: PqN, _ r: PqN) -> Bool {
            return l.p == r.p
        }

        public static func < (_ l: PqN, _ r:PqN) -> Bool {
            return l.p < r.p
        }
    }
    
    var pq = [PqN<E>]()

    var count:Int {
        pq.count
    }

    func insert(_ p:Int, _ o:E) {
        pq.append(PqN(p, o))
        shiftUp(pq.count - 1)
    }

    func pop() -> E? {
        if pq.count == 0 {
            return nil
        }
        let ret = pq[0]
        pq.swapAt(0, pq.count - 1)
        pq.removeLast()
        shiftDown(0)
        return ret.v
    }

    func shiftDown(_ idx:Int) {
        if idx == pq.count {
            return
        }
        let leftChild = (idx << 1) | 1
        let rightChild = (idx << 1) + 2
        if rightChild < pq.count {
            if pq[leftChild] <= pq[rightChild] && pq[idx] > pq[leftChild] {
                pq.swapAt(leftChild, idx)
                shiftDown(leftChild)
            } else if pq[leftChild] > pq[rightChild] && pq[idx] > pq[rightChild] {
                pq.swapAt(rightChild, idx)
                shiftDown(rightChild)
            }
        } else if leftChild < pq.count {
            if pq[idx] > pq[leftChild] {
                pq.swapAt(leftChild, idx)
                shiftDown(leftChild)
            }
        }
    }

    func shiftUp(_ idx: Int) {
        if idx == 0 {
            return
        }
        let p = (idx - 1) >> 1
        if pq[p] > pq[idx] {
            pq.swapAt(p, idx)
            shiftUp(p)
        }
    }
}