より良い実装

koherさんの実装のほうが良いのでコンテストで使う場合はこちらをオススメします。

優先度付きキュー(Priority Queue)

優先度付きキューは、取り出す先頭の要素が常に優先度最大(または最小)になるように内部で要素に順序をつけます。順序と言っても配列のように先頭から順に配置するのではなく、二分ヒープを使って順序性を実現します。

値の追加に $ O(\log N) $

値の取り出しに $ O(\log N) $ かかります。

一次元配列で二分ヒープを表現します。詳しい仕組みは下記の記事が大変わかりやすいです。

実装は以下を参考にしました。

以下実装です。

struct PriorityQueue<T> {
    private var data: [T]
    private var ordered: (T, T) -> Bool
    
    public var isEmpty: Bool {
        return data.isEmpty
    }
    
    public var count: Int {
        return data.count
    }

    init(_ order: @escaping (T, T) -> Bool) {
        self.data = []
        self.ordered = order
    }

    init<Seq: Sequence>(_ seq: Seq, _ order: @escaping (T, T) -> Bool) where Seq.Element == T {
        self.data = []
        self.ordered = order
        
        for x in seq {
            push(x)
        }
    }
    
    public mutating func pop() -> T? {
        return data.popLast().map { item in
            var item = item
            if !isEmpty {
                swap(&item, &data[0])
                siftDown()
            }
            return item
        }
    }
    
    public mutating func push(_ item: T) {
        let oldLen = count
        data.append(item)
        siftUp(oldLen)
    }
    
    private mutating func siftDown() {
        var pos = 0
        let end = count
        
        data.withUnsafeMutableBufferPointer { bufferPointer in
            let _data = bufferPointer.baseAddress!
            swap(&_data[0], &_data[end])
            
            var child = 2 * pos + 1
            while child < end {
                let right = child + 1
                if right < end && ordered(_data[right], _data[child]) {
                    child = right
                }
                swap(&_data[pos], &_data[child])
                pos = child
                child = 2 * pos + 1
            }
        }
        siftUp(pos)
    }
    
    private mutating func siftUp(_ pos: Int) {
        var pos = pos
        while pos > 0 {
            let parent = (pos - 1) / 2;
            if ordered(data[parent], data[pos]) {
                break
            }
            data.swapAt(pos, parent)
            pos = parent
        }
    }
}

extension PriorityQueue: Sequence, IteratorProtocol {
    mutating func next() -> T? {
        return pop()
    }
}

(shiftな気がしますがRustでの実装がsiftなのでsiftにしました。)

実装の概要

初期化

初期化にはシーケンス(任意)と比較方法を渡します。シーケンスを渡せば、各要素をヒープに格納します。

任意の比較方法

任意の比較方法を初期化時に指定できます。>=とすればmax-heap、<=とすればmin-heapになります。

シーケンス

SequenceとIteratorProtocolを実装しているため、for文で回すことができます。

var q = PriorityQueue([654,489,3,45,46,32,486,123], <=)

for x in q {
    print(x)
}

// 3
// 32
// 45
// 46
// 123
// 486
// 489
// 654

値取り出し(pop)の流れ

popでは先頭の要素を取り出しますが、そのまま取り出すと計算量が$ O(N) $になるため、取り出す際は、先頭と末尾の要素を入れ替え、popLastします。

こうすることで取り出し作業自体は $ O(1) $ になります。

[5, 4, 3, 2, 1] -swap-> [1, 4, 3, 2, 5] -popLast-> [1, 4, 3, 2], 5

そして先頭に来た末尾だった要素を正しい場所に持っていくため、二分ヒープの下へ下へ持っていきます(siftDownメソッド)。この処理に $ O(\log N) $ かかります。

    public mutating func pop() -> T? {
        return data.popLast().map { item in
            var item = item
            if !isEmpty {
                swap(&item, &data[0])
                siftDown()
            }
            return item
        }
    }

値追加(push)の流れ

pushではまず末尾に値を追加します。

[5, 4, 2, 1] <-push- 3
↓
[5, 4, 2, 1, 3]

そして、末尾の要素を正しい場所に持っていくため、二分ヒープの上へ上へ持っていきます(siftUpメソッド)。この処理に $ O(\log N) $ かかります。

    public mutating func append(_ item: T) {
        let oldLen = count
        data.append(item)
        siftUp(oldLen)
    }

使用例