【深さ優先探索】DFSによる全探索を理解したい！（Atcoder ABC119 C解説）【Python】

アルゴリズム

2022.09.03

はじめに

問題はこちら。

AtcoderのABC119 C – Synthetic Kadomatsuは、深さ優先探索（DFS）または4進数のbit全探索で解ける問題である。公式解説はPythonの正解例を紹介しているが、2022年9月現在茶コーダーの私には難しすぎたので、理解するのにかなり苦労した。そこで、本記事では、「どうすれば公式解説通りの実装に至るのか」、私の考え方を紹介する。

1. 深さ優先探索（dfs関数をどう考えるか）

過去の情報を次々と継承していく「フィボナッチ数列」的なf(n) = f(n-1) …のタイプではなく、f(n) = f(n+1) …のように、一手先の情報を参照する再帰関数をイメージできない！という問題を解決するのが本記事の最大の目的である。

よって、再帰関数自体を知らない方には、ネット上の入門的な記事をご覧いただきたい。

関数dfs(v, a, b, c)を、「始状態(v, a, b, c)からスタートし、終状態(A, B, C)に到達するまでにかかる最小コスト」であると考える。ここで、a, b, cは最終的に高さA, B, Cにしたい３本の竹の現在の高さであり、vはlのv番目 = l[v]（をどこに追加するか）を考えていることを表す。vは、深さ優先探索における深さとも言える。

図で表すと次のようにかける。lの１つ目の要素l[0]を竹Aに、２つ目の要素l[1]を竹Bに追加した場合は、dfs(v, a, b, c) = dfs(2, l[0], l[1], 0)となる。

また、dfsについての漸化式は次のように表すことができる。

\[
\mathrm{dfs}(v, a, b, c) =
\begin{cases}
\mathrm{min}\{\\
\mathrm{dfs}(v+1, a, b, c), \\
\mathrm{dfs}(v+1, \ a+l[v], \ b, \ c)+10, \\
\mathrm{dfs}(v+1, \ a, \ b+l[v], \ c)+10, \\
\mathrm{dfs}(v+1, \ a, \ b, \ c+l[v])+10\\
\} \ \ (v < N)\\
\\
|A-a| + |B-b| + |C-c| \ \ (v=N)
\end{cases}
\]

この漸化式で戸惑ったところは、始状態（a, b, c = 0の初期状態）は「v=0から始まり、vが深く（大きく）なるにつれてコストが加算されてゆくはず」なのに、dfs(v=N, …)が決まるとdfs(v=N-1, …)が決まり、…dfs(v=1, …)が決まるとdfs(0, 0, 0, 0)が決まることである。

これを理解するためには、上で述べた関数dfs()の定義をしっかりと頭に入れておく必要がある。関数dfs()は、深さv+1からスタートしたほうが、深さvからスタートするより少ないコストを返す関数であると考えることである。

考えてみれば当たり前で、N=8の時、dfs(v=7,…)（最大７本の竹を用いることのできる状態）から終状態を目指す方が、dfs（v=6）（最大６本の竹を用いることのできる状態）よりも、よりコストをかけずに終状態を目指すことができる。（竹をすべて使う必要は無いのだから）

2. Pythonでの実装例

N, A, B, C = map(int, input().split())
l = [int(input()) for _ in range(N)]

INF = 10**9
# 最低コストを出力する関数
# a, b, cは3本の竹の現在の長さ
# 現在l[v]の何番目を考えているか
def dfs(v, a, b, c):
    # N本目までの調査が終わった場合
    if v == N:
        # lの最後の要素まで探索して、全ての竹に高さがある場合
        # 最初、高さ0 => 1本竹を追加する操作は、「合成」ではない。
        # よってMP = 10 * 3本分の余計なコストを引く
        if min(a, b, c) > 0:
            return abs(A - a) + abs(B - b) + abs(C - c) - 30  
        # lの最後の要素まで探索して、高さ0の竹がある場合
        else:
            return INF
    # メイン部分
    ret0 = dfs(v + 1, a, b, c)
    ret1 = dfs(v + 1, a + l[v], b, c) + 10
    ret2 = dfs(v + 1, a, b + l[v], c) + 10
    ret3 = dfs(v + 1, a, b, c + l[v]) + 10
    return min(ret0, ret1, ret2, ret3)

print(dfs(0, 0, 0, 0))

最後に

理解するのに５時間以上かかったが、再帰関数の扱い方を頭の中でゴロゴロ転がしたり、深さ優先探索の概念と戯れたおかげで、少し賢くなった気がする。ちなみに、気分転換にランニングしながら考えていた時に理解できたので、身体を動かすのは大事だと思った。