Swiftには標準ライブラリに優先度付きキュー(Priority Queue)が無いので、自分で実装する必要があります。
この記事は競技プログラミングで使える優先度付きキューのSwiftでの私の実装を紹介します。
より良い実装
koherさんの実装のほうが良いのでコンテストで使う場合はこちらをオススメします。
優先度付きキュー(Priority Queue)
優先度付きキューは、取り出す先頭の要素が常に優先度最大(または最小)になるように内部で要素に順序をつけます。順序と言っても配列のように先頭から順に配置するのではなく、二分ヒープを使って順序性を実現します。
値の追加に $ O(\log N) $
値の取り出しに $ O(\log N) $ かかります。
一次元配列で二分ヒープを表現します。詳しい仕組みは下記の記事が大変わかりやすいです。
- https://medium.com/@yasufumy/data-structure-heap-ecfd0989e5be
- https://ufcpp.net/study/algorithm/col_heap.html
実装は以下を参考にしました。
- https://github.com/davecom/SwiftPriorityQueue/blob/master/Sources/SwiftPriorityQueue/SwiftPriorityQueue.swift
- https://github.com/raywenderlich/swift-algorithm-club/blob/master/Heap/Heap.swift
- https://doc.rust-lang.org/std/collections/binary_heap/struct.BinaryHeap.html
以下実装です。
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)
}
使用例