E - Multiple-Free Sequences 解説 by cirno3153

軽実装方針

今回の問題は \((x, y)\) としてあり得るものが \(O(M^2)\) 通りしかなく、また漸化式の形をしていることから、メモ化再帰で解く方法が有効です。

メモ化再帰で解くことを考えます。
初めにメモ用の配列の状態として、 \(M\) の倍数が含まれたことを示すNG\(M\) の倍数を含まなかったことを示すOK、未計算であることを示すUNKNOWNの3状態を用意します。 そして、メモ化配列を全てUNKNOWNで埋めておきます。

enum class State {
    OK, NG, UNKNOWN
}
class Solver(val M: Int, val A: Int, val B: Int) {
    val memo = Array(M) { Array(M) { State.UNKNOWN } }
}

次にメモ化再帰を行う関数solve(x, y)を用意します。 これは、数列の末尾2項が \((\ldots, x, y)\) だった時の答えを求める関数です。
まず、自明にmemo[x][y]UNKNOWNではなかった場合、その値を返します。
UNKNOWNだった場合、 \(x=0\) もしくは \(y=0\) ならばNG、そうでなければ \((\ldots, y, Ay + Bx \bmod M)\) の解に等しいです。
これをそのまま解くと無限ループする可能性がありますが、無限ループした場合はまさにOKということになるので、予めmemo[x][y]OKを入れておくと無限ループを回避できます。

fun solve(x: Int, y: Int): State {
    if (memo[x][y] == State.UNKNOWN) {
        memo[x][y] = State.OK
        val next = (A * y + B * x) % M
        memo[x][y] = if (x == 0 || y == 0) State.NG else solve(y, next)
    }
    return memo[x][y]
}

以上を纏めると、以下のようなコードになります。

enum class State {
    OK, NG, UNKNOWN
}
class Solver(val M: Int, val A: Int, val B: Int) {
    val memo = Array(M) { Array(M) { State.UNKNOWN } }
    fun solve(x: Int, y: Int): State {
        if (memo[x][y] == State.UNKNOWN) {
            memo[x][y] = State.OK
            val next = (A * y + B * x) % M
            memo[x][y] = if (x == 0 || y == 0) State.NG else solve(y, next)
        }
        return memo[x][y]
    }
    fun ans(): Int {
        return (0 until M).sumOf { i ->
            (0 until M).count { j -> solve(i, j) == State.OK }
        }
    }
}
fun main() {
    val (M, A, B) = readLine()!!.split(" ").map { it.toInt() }
    println(Solver(M, A, B).ans())
}

計算量を求めます。
メモ化配列に値を代入する回数を考えると、UNKNOWNが入っている時のみそれ以外の値を \(O(1)\) 回代入するので配列長と同じ \(O(M^2)\) となります。

このように、漸化式など直前の状態から計算可能な形で与えられる問題で、全状態を管理しても空間計算量が小さいものにはメモ化再帰が実装を軽くするテクニックとして有効です。
グラフ上のDFSや木DP、フラクタルな構造をしている問題にも使える場合が多いので、まずメモ化再帰を検討してみるのも良いかもしれません。

投稿日時:
最終更新: