AtCoder B,C問題をKotlinで解こう - ABC253

当記事ではAtCoder、ABCのB問題ならびにC問題（時々D問題も）のKotlinでの解法を超初心者向けに詳細に解説します。

B - Distance Between Tokens

やりたいこと

oのマス2つ（それぞれa,bとする）の座標の マンハッタン距離 を求めよう。マンハッタン距離は Math.abs(a.x - b.x) + Math.abs(a.y - b.y) となる。

入力値の取得

    // 入力値の取得
    val (h, w) = readLine()!!.split(" ").map { it.toInt() }
    val s = (1..h).map { readLine()!! }

座標保持用のデータクラス作成

// 座標保持用のデータクラス
data class Pos(val y: Int, val x: Int)

o のマスの座標取得

    // oのマスの座標取得
    val token = mutableListOf<Pos>()
    for (i in s.indices) {
        for (j in s[i].indices) {
            if (s[i][j] == 'o') {
                // oのマスが見つかったので、データクラスを作成し保持する
                token.add(Pos(i, j))
            }
        }
    }
    // oのマスは2つあることが保証されているのでtokenのサイズは2となる。

サンプルコード

main.kt

fun main(args: Array<String>) {
    // 入力値の取得
    val (h, w) = readLine()!!.split(" ").map { it.toInt() }
    val s = (1..h).map { readLine()!! }

    // oのマスの座標取得
    val token = mutableListOf<Pos>()
    for (i in s.indices) {
        for (j in s[i].indices) {
            if (s[i][j] == 'o') {
                // oのマスが見つかったので、データクラスを作成し保持する
                token.add(Pos(i, j))
            }
        }
    }
    // oのマスは2つあることが保証されているのでtokenのサイズは2となる。
    println(Math.abs(token[0].y - token[1].y) + Math.abs(token[0].x - token[1].x))
}

// 座標保持用のデータクラス
data class Pos(val y: Int, val x: Int)

C - Max - Min Query

やりたいこと

sは入力されるxをキーとしたmapにし、値で要素数を管理しよう。
sのキーの最大と最小をすぐ取り出せるよう、最大管理用のPriorityQueueと最小管理用のPriorityQueueを用意しよう。

PriorityQueueは昇順で最初になるものが先頭に来ることを保証しますが、その他の順序は全く保証しません。末尾のものが昇順で最後の要素ということにはならないため、最大を取り出すためには最大を管理するためのPriorityQueueが別途必要となります。

入力値の取得

    // 入力値の取得
    val q = readLine()!!.toInt()

    repeat(q) { _ ->
        val query = readLine()!!.split(" ").map { it.toInt() }
        when (query[0]) {
            1 -> {
                // 追加クエリ
                val x = query[1]
            }
            2 -> {
                // 削除クエリ
                val (x,c) = query.drop(1)
            }
            else -> {
                // 出力クエリ
            }
        }
    }

sへの値の追加と削除

    val s = mutableMapOf<Int, Int>()
    repeat(q) { _ ->
        val query = readLine()!!.split(" ").map { it.toInt() }
        when (query[0]) {
            1 -> {
                // 追加クエリ
                val x = query[1]
                s[x] = (s[x] ?: 0) + 1
            }
            2 -> {
                // 削除クエリ
                val (x,c) = query.drop(1)
                val remove = Math.min(s[x] ?: 0, c)
                s[x] = (s[x] ?: 0) - remove
            }
            else -> {
                // 出力クエリ
            }
        }
    }

sのキーの最大と最小の管理

sのキーの最大と最小を管理するためのPriorityQueuueを二つ作成しよう。

    // 最小管理用PriorityQueue
    val keysMin = PriorityQueue<Int>()
    // 最大管理用のPriorityQueue
    val keysMax = PriorityQueue<Int>()

追加クエリの処理時に、追加前の状態で該当要素数が0の場合、それぞれのPriorityQueueに要素を追加しよう。

            1 -> {
                // 追加クエリ
                val x = query[1]
                if ((s[x] ?: 0) <= 0) {
                    keysMin.add(x)
                    keysMax.add(x * -1)
                }
                s[x] = (s[x] ?: 0) + 1
            }

PriorityQueueは常に昇順のものを先頭に持ってきます。整数を要素にするときは最小が先頭になるため、最大を先頭にしたいときは、PriorityQueueへの追加と取り出し時に -1 をかける等の処理が必要になります。

最大 - 最小 の取得

                // 出力クエリ
                
                // 削除クエリにより要素数がゼロになっているものを削除する。
                while (keysMin.isNotEmpty() && keysMin.peek().let { s[it] ?: 0 } <= 0) {
                    keysMin.poll()
                }
                while (keysMax.isNotEmpty() && keysMax.peek().let { s[it * -1] ?: 0 } <= 0) {
                    keysMax.poll()
                }
                
                // keyMaxに要素を追加するときに -1 をかけているので、取り出し時にも -1 をかける
                keysMax.peek() * -1 - keysMin.peek()

サンプルコード

main.kt

import java.lang.StringBuilder
import java.util.*

fun main(args: Array<String>) {
    // 入力値の取得
    val q = readLine()!!.toInt()
    val s = mutableMapOf<Int, Int>()
    // 最小管理用PriorityQueue
    val keysMin = PriorityQueue<Int>()
    // 最大管理用のPriorityQueue
    val keysMax = PriorityQueue<Int>()
    val ans = StringBuilder()
    repeat(q) { _ ->
        val query = readLine()!!.split(" ").map { it.toInt() }
        when (query[0]) {
            1 -> {
                // 追加クエリ
                val x = query[1]
                if ((s[x] ?: 0) <= 0) {
                    keysMin.add(x)
                    keysMax.add(x * -1)
                }
                s[x] = (s[x] ?: 0) + 1
            }
            2 -> {
                // 削除クエリ
                val (x,c) = query.drop(1)
                val remove = Math.min(s[x] ?: 0, c)
                s[x] = (s[x] ?: 0) - remove
            }
            else -> {
                // 出力クエリ

                // 削除クエリにより要素数がゼロになっているものを削除する。
                while (keysMin.isNotEmpty() && keysMin.peek().let { s[it] ?: 0 } <= 0) {
                    keysMin.poll()
                }
                while (keysMax.isNotEmpty() && keysMax.peek().let { s[it * -1] ?: 0 } <= 0) {
                    keysMax.poll()
                }
                // keyMaxに要素を追加するときに -1 をかけているので、取り出し時にも -1 をかける
                ans.appendln(keysMax.peek() * -1 - keysMin.peek())
            }
        }
    }
    print(ans.toString())
}