Open2

Leetcode自分用解答メモ[Explore - Heap]

リアルにサーモンリアルにサーモン

https://leetcode.com/explore/learn/card/heap/646/practices/4014/

object Solution {
    import scala.collection.mutable.PriorityQueue
    def findKthLargest(nums: Array[Int], k: Int): Int = {
        val N = nums.length
        val pq = PriorityQueue[Int]()(Ordering.by((i: Int) => i).reverse)
        
        for (i <- 0 until k) pq.enqueue(nums(i))
        for (i <- k until N) {
            val num = nums(i)
            if (pq.head < num) {
                pq.dequeue
                pq.enqueue(num)
            }
        }
        
        pq.head
    }
}

全ての要素をヒープに入れてから取り出す実装(O(N + klogN)時間、O(N)スペース)では単純なので、

この実装だと、O(Nlogk)時間、O(k)スペースで解くことができる。

リアルにサーモンリアルにサーモン

https://leetcode.com/problems/kth-largest-element-in-a-stream/

ストリーム形式で入力される数字からk番目に大きい数を取り出すデータ構造を作る。

class KthLargest(_k: Int, _nums: Array[Int]) {
    import scala.collection.mutable
    
    val k = _k
    val pq = mutable.PriorityQueue[Int]()(Ordering.by(i => (i: Int))).reverse
    
    val N = _nums.length
    for (i <- 0 until k.min(N)) pq.enqueue(_nums(i))
    for (i <- k until N) add(_nums(i))
    
    def add(`val`: Int): Int = {
        if (pq.size < k) pq.enqueue(`val`)
        else if (pq.head < `val`) {
            pq.dequeue
            pq.enqueue(`val`)
        }
        pq.head
    }
}

/**
 * Your KthLargest object will be instantiated and called as such:
 * var obj = new KthLargest(k, nums)
 * var param_1 = obj.add(`val`)
 */

初期化時にk - 1個の要素しか与えられていない場合のみ注意が必要。