優先度付きキュー

読了目安約3分

二分ヒープを自作し、最小値の追加と取り出しを O(log N) で行う。

この章の目次

優先度付きキューは、追加した要素の中の最小値 (または最大値) を素早く取り出せる入れ物です。配列で最小値を探すと 1 回 O(N) かかりますが、優先度付きキューなら追加も取り出しも O(log N) で済みます。

Swift の標準ライブラリには優先度付きキューがありません。そこで、実現方法の定番である二分ヒープを自作します。

二分ヒープの形

二分ヒープは、配列で表した木です。 items[i] の子を items[2 * i + 1] と items[2 * i + 2]、親を items[(i - 1) / 2] と決めます。

テキスト
        1        ← items[0]
      /   \
     3     2     ← items[1], items[2]
    / \   /
   7   5 4       ← items[3], items[4], items[5]

この木の実体は、配列 [1, 3, 2, 7, 5, 4] だけです。そのうえで「親 ≦ 子」というルールを常に保ちます。すると根 items[0] が必ず全体の最小値になります。

push: 末尾に足して上げる

追加する値は、まず配列の末尾に置きます。そのあと親と比べ、親より小さい間は交換しながら上がります (up-heap)。

上の図に 0 を push すると、4 の右隣に入り、親の 2 と交換、さらに親の 1 と交換して根に着きます。 N 個の要素を持つ木の高さは log N 程度なので、交換は O(log N) 回で止まります。

pop: 末尾を根に移して下げる

最小値は根にあるので、取り出すのは items[0] です。空いた根には末尾の要素を移し、今度は 2 つの子の小さい方と比べ、子より大きい間は交換しながら下がります (down-heap)。こちらも O(log N) 回で止まります。

struct Heap を書き切る

Swift

struct Heap {
    private var items: [Int] = []

    var isEmpty: Bool { items.isEmpty }
    var count: Int { items.count }

    mutating func push(_ value: Int) {
        items.append(value)
        var i = items.count - 1
        while i > 0 {
            let parent = (i - 1) / 2
            if items[parent] <= items[i] { break }
            items.swapAt(parent, i)
            i = parent
        }
    }

    mutating func pop() -> Int? {
        guard let top = items.first else { return nil }
        items.swapAt(0, items.count - 1)
        items.removeLast()
        var i = 0
        while true {
            let left = 2 * i + 1
            let right = 2 * i + 2
            var smallest = i
            if left < items.count && items[left] < items[smallest] { smallest = left }
            if right < items.count && items[right] < items[smallest] { smallest = right }
            if smallest == i { break }
            items.swapAt(i, smallest)
            i = smallest
        }
        return top
    }
}

var heap = Heap()
for x in [5, 2, 8, 1, 9] {
    heap.push(x)
}
while let value = heap.pop() {
    print(value)
}

どんな順で push しても、pop は必ず小さい順に返します。

試してみよう: push する数列を変えて、出力が常に昇順になることを確かめてください。

用途: 上位 K 個

「N 個のうち大きい方から K 個」は、大きさ K のヒープで取れます。

Swift
var heap = Heap()
for x in numbers {
    heap.push(x)
    if heap.count > k {
        _ = heap.pop()  // 一番小さいものから捨てる
    }
}

ヒープには常に「ここまでの大きい方 k 個」だけが残ります。全体をソートする O(N log N) に対し、この方法は O(N log K) です。

最大値側を取り出したいときは、符号を反転した値を push し、pop の結果をまた反転するのが手軽です。

学びどころ

概念	一言まとめ
優先度付きキュー	最小値の追加と取り出しが速い入れ物
二分ヒープ	「親 ≦ 子」を保つ、配列で表した木
push	末尾に足して up-heap。O(log N)
pop	根を返し、末尾を根に移して down-heap。O(log N)

演習

1 行目にクエリ数 Q が入力されます。続く Q 行は、次のどちらかです。

1 x：整数 x を追加する
2：最小値を取り出して出力する (このとき要素は必ず 1 個以上あります)

スターターの Heap を使って処理してください。

テキスト
入力例:
6
1 5
1 2
2
1 1
2
2

出力例:
2
1
5

模範解答

Swift
struct Heap {
    private var items: [Int] = []

    var isEmpty: Bool { items.isEmpty }
    var count: Int { items.count }

    mutating func push(_ value: Int) {
        items.append(value)
        var i = items.count - 1
        while i > 0 {
            let parent = (i - 1) / 2
            if items[parent] <= items[i] { break }
            items.swapAt(parent, i)
            i = parent
        }
    }

    mutating func pop() -> Int? {
        guard let top = items.first else { return nil }
        items.swapAt(0, items.count - 1)
        items.removeLast()
        var i = 0
        while true {
            let left = 2 * i + 1
            let right = 2 * i + 2
            var smallest = i
            if left < items.count && items[left] < items[smallest] { smallest = left }
            if right < items.count && items[right] < items[smallest] { smallest = right }
            if smallest == i { break }
            items.swapAt(i, smallest)
            i = smallest
        }
        return top
    }
}

let q = Int(readLine()!)!
var heap = Heap()
for _ in 0..<q {
    let query = readLine()!.split(separator: " ").map { Int($0)! }
    if query[0] == 1 {
        heap.push(query[1])
    } else {
        print(heap.pop()!)
    }
}