08. Kth Largest Element in a Stream

読了目安約2分

Heap / LeetCode 703 (Easy) — ストリームに値を追加しながら K 番目に大きい値を返すクラスを設計する。

この章の目次

問題

LeetCode 703. Kth Largest Element in a Stream（Easy / カテゴリ: Heap）

次々に値が届くストリームから、K 番目に大きい値を答え続けるクラスを設計する問題です。 init(k, nums) で K と初期値の配列を受け取ります。 add(val) は値を 1 つ追加し、その時点で K 番目に大きい値を返します。add は繰り返し呼ばれます。

考え方

ヒント 1

add のたびに全体をソートし直しても解けますが、毎回 O(n log n) かかります。「K 番目に大きい値」だけが要るのだから、大きい方から K 個だけを保てば十分です。

ヒント 2

大きい方から K 個を最小ヒープ（最小値を O(log n) で取り出せる二分ヒープ）で持ちます。その最小値、つまりヒープの先頭がちょうど K 番目に大きい値です。追加してサイズが K を超えたら、最小値を捨てます。

Swift 実装のポイント

Swift の標準ライブラリに優先度付きキューはありません。配列を使った二分ヒープを自分で実装します。
ヒープは struct + mutating func で書けます。親は (i - 1) / 2、子は 2 * i + 1 と 2 * i + 2 です。
init の中でも add を使い回すと、サイズを K 以下に保つ処理が 1 か所にまとまります。戻り値は _ = で捨てます。

模範解答

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

    var count: Int { items.count }
    var top: Int? { items.first }

    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() {
        guard !items.isEmpty else { return }
        items[0] = items[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
        }
    }
}

class KthLargest {
    private var heap = MinHeap()
    private let k: Int

    init(_ k: Int, _ nums: [Int]) {
        self.k = k
        for num in nums {
            _ = add(num)
        }
    }

    func add(_ val: Int) -> Int {
        heap.push(val)
        if heap.count > k {
            heap.pop()
        }
        return heap.top!
    }
}

計算量: add 1 回あたり O(log K) 時間、O(K) 追加メモリ。