【AtCoder】ABC234の「D - Prefix K-th Max」をScalaのPriorityQueue使用して解く

概要

AtCoder Beginner Contest 234のD問題は、こちらの解説にある通り、優先度付きキューを使用して解くことができます。
Scalaでも優先度付きキューの機能はありますが、あまり情報がなかったので、今回は解いた時のメモ書きを残しておきます。

対応方針

• 優先度付きキューにOrdering[Int].reverseを設定する。これによって、値の小さいものが上に来る。
• ループの項番がKの数になるまで、優先度付きキューに値を追加していく。
• ループの項番がKの数になったら、結果として先頭の値を取得する。
• ループの項番がKより大きい数の場合は、優先度付きキューに追加後にdequeueで先頭を削除した後に、結果として先頭の値を取得する。
• PriorityQueueの仕様の詳細は、こちらのドキュメントを参照。

実装

以下のコードの提出結果はこちらです。

object Main extends App {
  // nとkの入力
  val nkVector = scala.io.StdIn.readLine().split(" ").map(_.toInt).toVector
  val n = nkVector(0)
  val k = nkVector(1)
  // pの配列を入力
  val pArray = scala.io.StdIn.readLine().split(" ").map(_.toInt)
  // 結果の配列定義
  var results = new Array[Int](n - k + 1)
  // 優先度付きキューの定義
  implicit val orderP = Ordering[Int].reverse
  val pq = new collection.mutable.PriorityQueue()
  // キューに追加していく処理
  pArray.zipWithIndex.foreach(p => {
    // キューに追加
    pq.addOne(p._1)
    // indexがkの時はキューの先頭を結果に入れる
    if (p._2 == k - 1) {
      // 先頭を結果に追加
      results(p._2 - (k - 1)) = pq.head
    // indexがkより大きい時は、キューの先頭を削除した上で先頭を結果に入れる
    } else if (p._2 > k - 1) {
      // キューの先頭を削除
      pq.dequeue()
      // 先頭を結果に追加
      results(p._2 - (k - 1)) = pq.head
    }
  })
  // 結果出力
  println(results.mkString(sep = System.lineSeparator()))
}