👻

ABC273 D - LRUD Instructions のメモ

2023/10/19に公開

重実装問題を募集したら教えてもらったので解きました。解けたのでよかった。
https://atcoder.jp/contests/abc273/tasks/abc273_d

考察

クエリ一つ一つに対して1マスしか移動しないのであれば愚直にやればいいのですが、実際には最大10^9回移動するので、計算で求めるしかありません。一直線に進むので移動先のマスがどこなのかは簡単に求められますが、壁があるとそれ以上進めないので、それによって問題が難しくなっています。壁の有無と、ある場合は壁の位置を高速に求めることができれば解けるということですね。

一度に移動するマス数は制約が大きいのでそちらを探索することはできない(事前にデータ構造を作っておくとかも無理)ですが、壁の数は2×10^5以下なので、いかにもこちらを探索すればよさそうな感じです。
高速な探索というと二分探索が思い浮かびます。なんとか二分探索できないかなあと思ったのですが、できるってことに気づいたのでそれで進めました。

最近見たいくつかの問題で、Kotlinの二分探索メソッドを使うと要素が存在する場合の位置だけでなく、要素が存在しない場合であっても、仮に存在した場合の位置を知ることができるという知識を得ていまいした。
https://zenn.dev/dhirabayashi/articles/3a8055d3f413be

これが使えるんじゃないかと。グリッドではなく、壁の一覧を列挙してソートしただけのものに対して二分探索する想定です。
それで、移動元のマスと移動先のマスの両方で二分探索すると、その間に壁が存在しないのであれば戻り値が一致するはずということにまず気づきました。じゃあ壁が存在する場合はどうなるのかというと、その場合は移動元のほうで探索した結果が示す位置の隣が最初の壁です。(グリッドの隣ではなく、壁の配列の要素の隣)
移動元自体は明らかに壁ではないので、移動元での探索がヒットすることはなく、戻り値は確実に要素が存在しない場合の値です。それは「仮に存在した場合の位置」なので、最初の壁の一つ前の要素を指すはずなので。

これがわかれば、あとはその壁の配列をどうやって作るかです。二分探索を使うためにはソートする必要があります。右へ移動するケースについて考えると、高さが第一条件、幅が第二条件として昇順ソートすればいいとわかりました。左の場合は、左右が逆なので幅は降順ソートになりますが、高さについては同様に第一条件として昇順ソートとすればいいですね。上下の移動については幅と高さのソート条件の順番を変えればいけます。
どのように実装に落とし込むかですが、4方向しかないので配列とComparatorを4つずつ作ってしまうことで対処しました。

あとは移動先を求めて(移動先がグリッドの範囲外にならないようにして)、二分探索で壁の有無と位置を求めて、それをもとに答えとなる位置を求めて出力、をクエリ分実行すればいいです。

実装

まあ実装が重めの問題なので、ここからどう実装するのかのほうが大変だったかもしれません。
混乱しそうだったのでコメントを書きながら実装していきました。

import java.io.PrintWriter
import java.lang.RuntimeException
import kotlin.math.max
import kotlin.math.min

fun main() {
    val br = System.`in`.bufferedReader()
    val (h, w, rs, cs) = br.readLine().split(" ").map { it.toInt() }
    val n = br.readLine().toInt()
    val rc = Array(n) {
        val (r, c) = br.readLine().split(" ").map { it.toInt() }
        Point(r, c)
    }
    val q = br.readLine().toInt()
    val dl = Array(q) {
        val (d, l) = br.readLine().split(" ")
        d.first() to l.toInt()
    }

    // 二分探索用にソート
    val leftComparator = compareBy<Point> { it.r }.thenByDescending { it.c }
    val leftRc = rc.sortedWith(leftComparator)

    val rightComparator = compareBy<Point> { it.r }.thenBy { it.c }
    val rightRc = rc.sortedWith(rightComparator)

    val upComparator = compareBy<Point> { it.c }.thenByDescending { it.r }
    val upRc = rc.sortedWith(upComparator)

    val downComparator = compareBy<Point> { it.c }.thenBy { it.r }
    val downRc = rc.sortedWith(downComparator)

    val comparatorMap = mapOf(
        'L' to (leftComparator to leftRc),
        'R' to (rightComparator to rightRc),
        'U' to (upComparator to upRc),
        'D' to (downComparator to downRc),
    )

    val diffMap = mapOf(
        'L' to (0 to 1),
        'R' to (0 to -1),
        'U' to (1 to 0),
        'D' to (-1 to 0),
    )

    // メイン処理
    var currentR = rs
    var currentC = cs
    PrintWriter(System.out).use {
        for((d, l) in dl) {
            // 移動先
            val (toR, toC) = when(d) {
                'L' -> currentR to (max(1, currentC - l))
                'R' -> currentR to (min(w, currentC + l))
                'U' -> (max(1, currentR - l)) to currentC
                'D' -> (min(h, currentR + l)) to currentC
                else -> throw RuntimeException()
            }

            // 途中に壁がある場合は特定する
            val (comparator, walls) = comparatorMap[d]!!
            val start = Point(currentR, currentC)
            val end = Point(toR, toC)

            val startIndex = walls.binarySearch(start, comparator)
            val endIndex =walls.binarySearch(end, comparator)

            // 一致したら壁はない、一致しなかったらある
            if(startIndex != endIndex) {
                val wallIndex = -(startIndex + 1)
                val (wallR, wallC) = walls[wallIndex]
                val (dr, dc) = diffMap[d]!!

                currentR = wallR + dr
                currentC = wallC + dc
            } else {
                currentR = end.r
                currentC = end.c
            }

            it.println("$currentR $currentC")
        }
    }
}

data class Point(
    val r: Int,
    val c: Int,
)

https://atcoder.jp/contests/abc273/submissions/46720049

位置を表すPointクラスを作っています。Pairでもできるのですが、firstとsecondのどっちがどっちなのかよく混乱するのでクラスを作りました。Pointでソートしますが、Comparatorを使うのでこれ自体はComparableにしなくてOKでした。
Comparatorは事前ソートで使うほか、二分探索でもメソッドに渡します。

出力が多いのでPrintWriterを使います。今回はuseを使って自動closeするようにしてみました。

壁がある場合、壁の位置がわかったらその1つ前の位置を求める必要があります。それが答えなので。どの方向に進むのかによって変わりますが、どこに移動するのかはMapで管理しました。

Kotlinの良さのおかげですが、意外とすっきりしたコードとして実装できた気がします。
(執筆時間: 42分57秒)

Discussion