F - AtCoder Express 3 解説
by 
square1001
本コンテストの最終問題「AtCoder Express 3」です。
まず、\(2^q\) 通りの鉄道路線すべてに対して駅 \(0\) から駅 \(N\) までの最短経路を求めていると、計算量 \(O(2^q \times N)\) かかります。すると、? の個数が 1 増えるごとに計算時間が 2 倍になるので、\(q = 30\) くらいでも実行時間制限に間に合いません。
これを改善するために、動的計画法 (DP) による効率的な解法を考えます。しかし、駅 \(0\) から駅 \(N\) までの最短経路を 幅優先探索 (BFS) で求めるアルゴリズムをベースとして考えると、複雑な状態を管理する必要があるため、動的計画法に応用するのは難しいです。
そのため、幅優先探索を使わないもっと単純な方法で駅 \(0\) から駅 \(N\) までの最短経路を求めるアルゴリズムを考え、これをベースとして本問題を解く動的計画法のアルゴリズムを構築していきます。
解説の目次
本解説では、以下のステップで解説をします。
まずは本問題より少し簡単な、「番号が大きくなる方向にしか進めない」場合（一方通行パターン）の解法を考えます。
- ステップ 1：「一方通行パターン」における最短経路長の求め方
- ステップ 2：「一方通行パターン」の数え上げ問題を \(O(N^3)\) で解く
これをベースに、本問題の解法を以下のステップで作り上げていきます。
- ステップ 3：駅 \(0\) から駅 \(N\) に行く最短経路長の求め方
- ステップ 4：この問題の動的計画法による \(O(N^3)\) 解法
- ステップ 5：ひと工夫で計算量 \(O(N^2)\) に高速化
そして、6 節には「解法のまとめ」を書きました。その後の 7 節に、本問題の C++ による実装例（サンプルコード）も載せています。なお、急いで解説を読みたい人は、6 節「解法のまとめ」を読んでください。
また、解法をさらに工夫することで、理論的にはこの問題を計算量 \(O(N \log^2 N)\) で解くこともできます。これについても 8 節で紹介します。
ステップ 1：「一方通行パターン」における最短経路長
ステップ 1 では、駅 \(1\) から駅 \(N-1\) までの管理するグループがすべて決まっている場合に最短経路を求める問題を考えます。この問題を幅優先探索を使わずに、前から順に処理して答えを求めるアルゴリズムを考えてみましょう。特に、状態数を削るためにも「答えを求めるために使用する変数の情報量」を少なくすることを考えましょう。
駅 \(i\) \((1 \leq i \leq N-1)\) を管理するグループが文字 \(c_i\) (A＝光速派、B＝準急派) で与えられる。駅 \(0\) と駅 \(N\) は両方のグループが管理する。
駅の番号が大きくなる方向にしか進めないとき、最短何分で駅 \(0\) から駅 \(N\) に移動できるか求めてください。
以下の図は、\(N = 8\)、\(c = \) AABABBB の例を示しています。丸に書かれた数は駅 \(0\) からの最短経路を表しており、駅 \(N\) までの最短経路は \(3\) です。
上の図で示されるように右方向にしか進まないので、左から順に「駅 \(0\) から駅 \(i\) までの最短経路の長さ \(d_i\)」を求めていくことができます。
アルゴリズム #1 \(i = 1, 2, \dots, N-1\) の順に、\(d_i\) を以下のように求める。
- \(c_i\) が
Aの場合：\(a\) を直前の光速派の駅として、普通列車に乗って駅 \(i-1 \rightarrow i\) に移動する方法と、光速列車に乗って駅 \(a \rightarrow i\) に移動する方法があるので、\(d_i = \min(d_{i-1}, d_a) + 1\) となる。- \(c_i\) が
Bの場合：\(b\) を直前の準急派の駅として、普通列車に乗って駅 \(i-1 \rightarrow i\) に移動する方法と、準急列車に乗って駅 \(b \rightarrow i\) に移動する方法があるので、\(d_i = \min(d_{i-1}, d_b) + 1\) となる。最終的に、駅 \(N\) までの最短経路は \(d_N = \min(d_{N-1}, d_a, d_b) + 1\) で求められる。
このように、前から順に最短経路を計算するとき、\(d_{i-1}, d_a, d_b\) の値だけを使って \(d_i\) の値を計算することができるのです。なので、\(3\) つの値 \((p, \alpha, \beta) = (d_{i-1}, d_a, d_b)\) だけを記録する、以下のアルゴリズムで答えが求められます。
アルゴリズム #2 \(i = 1, 2, \dots, N-1\) の順に、以下のように \((p, \alpha, \beta)\) を更新していく。初期値は \((p, \alpha, \beta) = 0\) である。
- \(c_i\) が
Aの場合：\(d_i = \min(d_{i-1}, d_a) = \min(p, \alpha)\) なので、\(p\) と \(\alpha\) の値を \(\min(p, \alpha) + 1\) に更新する。- \(c_i\) が
Bの場合：\(d_i = \min(d_{i-1}, d_b) = \min(p, \beta)\) なので、\(p\) と \(\beta\) の値を \(\min(p, \beta) + 1\) に更新する。最終的に、駅 \(N\) までの最短経路は \(\min(p, \alpha, \beta) + 1\) で求められる。
要約すれば、「直前の駅までの最短経路長」「直前の光速派の駅までの最短経路長」「直前の準急派の駅までの最短経路長」の \(3\) つの値を記録するだけで最短経路長が計算できる、ということです。
さらには、直前が「光速派の駅」ならば \(p = \alpha\) に、「準急派の駅」ならば \(p = \beta\) となります。ですから、実際は「直前の光速派の駅までの最短経路長」「直前の準急派の駅までの最短経路長」の \(2\) つの値しか記録する必要はありません。アルゴリズム #2 にこの改善を入れると、以下のようなアルゴリズムになります。
アルゴリズム #3 \(i = 1, 2, \cdots, N-1\) の順に、以下のように \((\alpha, \beta)\) を更新していく。初期値は \((\alpha, \beta) = (0, 0)\) である。
- \(c_i\) が
A、\(c_{i-1}\) がAの場合：\(\alpha\) の値を \(\alpha + 1\) に更新する。- \(c_i\) が
A、\(c_{i-1}\) がBの場合：\(\alpha\) の値を \(\min(\alpha, \beta) + 1\) に更新する。- \(c_i\) が
B、\(c_{i-1}\) がAの場合：\(\beta\) の値を \(\min(\beta, \alpha) + 1\) に更新する。- \(c_i\) が
B、\(c_{i-1}\) がBの場合：\(\beta\) の値を \(\beta + 1\) に更新する。- ただし \(i = 1\) の場合は、\(c_{i-1}\) としてどちらを選んでも、次の \((\alpha, \beta)\) の値は同じになる。
最終的に、駅 \(N\) までの最短経路は \(\min(\alpha, \beta) + 1\) で求められる。
ここで、\(0 \leq \alpha \lt N, 0 \leq \beta \lt N\) なので、管理する状態数を \(N^2\) 通りに絞ることができました。
ステップ 2：「一方通行パターン」における数え上げ問題
ステップ 1 で行った考察をベースに、「一方通行パターン」の問題を解く動的計画法のアルゴリズムを設計しましょう。左の駅から順に処理していくアルゴリズム #3 をベースにすると、DP に必要な情報は、以下の通りです。
- 直前に見た駅の番号 \(pos\)
- 直前の光速派の駅までの最短経路長 \(\alpha\)
- 直前の準急派の駅までの最短経路長 \(\beta\)
- 直前の駅 \(pos\) を管理するグループ \(pre\)（0 ＝ 光速派、1 ＝ 準急派とする）
この状態が何通りあるかを \(dp[pos][\alpha][\beta][pre]\) で記録します。ここで、アルゴリズム #3 と同様に、状態 \((pos, \alpha, \beta, pre)\) からは、次のように状態が変わります。
\(pre = 0\)（直前が光速派）のとき
- \(c_{pos+1}\) が
Aまたは?のとき、状態を \((pos+1, \alpha + 1, \beta, 0)\) に変えられる- \(c_{pos+1}\) が
Bまたは?のとき、状態を \((pos+1, \alpha, \min(\beta, \alpha) + 1, 1)\) に変えられる\(pre = 1\)（直前が準急派）のとき
- \(c_{pos+1}\) が
Aまたは?のとき、状態を \((pos + 1, \min(\alpha, \beta) + 1, \beta, 0)\) に変えられる- \(c_{pos+1}\) が
Bまたは?のとき、状態を \((pos + 1, \alpha, \beta + 1, 1)\) に変えられる
このようにして \(dp[pos][\alpha][\beta][pre]\) の値を順々に求めていくことで、計算量 \(O(N^3)\) で「一方通行パターン」の問題を解くことができます。
なお、ステップ 5 で説明するアイデアを使うことで、計算量 \(O(N^2)\) に改善することもできます。解説を読み終わったら、ぜひ考えてみましょう。
ステップ 3：駅 \(0\) から駅 \(N\) に行く最短経路長の求め方
ここからは、一方通行とは限らず、逆向きに移動することが可能な場合を考えます。ここで、次のような疑問を思いうかべる人もいるかもしれません。
疑問 逆向きに移動しないものが最短経路になるのではないか？最短経路のうち 1 つは、全て順方向に移動するものなのではないか？
残念ながら、これは成り立ちません。例えば \(N = 8\)、\(c =\) AAAABBB の場合は、最短経路で移動する方法は駅 \(0 \rightarrow 5 \rightarrow 4 \rightarrow 8\) の 1 つしかありません。下図のように、順方向の移動に限定してしまうと、最短経路長が \(4\) になってしまいます。
しかし、実は逆向きに移動することを許容する場合でも、最短経路には次のような性質があります。
性質 最短経路において逆方向に移動しうるのは、普通列車を使うときだけである。これも、\(2\) 連続以上で逆方向に進んではならない。
理由 光速列車や準急列車を一度でも逆方向に移動したり、\(2\) 連続で普通列車を逆方向に進んだりすると、同じ駅で \(2\) 回以上停まることになるから。そして、同じ駅で \(2\) 回以上停まるような経路は、最短経路にはなりえないから。
以下の図のように鉄道路線のグラフを置き換えてみると分かりやすいと思います。赤色・青色の辺を逆向きに通ったり、灰色の辺を \(2\) 連続で通ったりすると、一番左から一番右に行くときに同じ頂点を \(2\) 回通ってしまいます。
別のとらえ方として、\(2 \times n\) のマス目の左上から右下まで同じマスを \(2\) 度通らずに移動するときに、左に移動することはない、というものもあります。
つまり、ほとんどステップ 1 と同様に解くことができるということです。光速電車（赤色の辺）と準急電車（青色の辺）は順方向の向きを付けることができます。また、灰色の辺を逆向きに \(2\) 連続で移動することもありません。
これより、アルゴリズム #3 を少し変えた以下のアルゴリズムで最短経路長を計算することができます。
アルゴリズム #4 \(i = 1, 2, \cdots, N-1\) の順に、以下のように \((\alpha, \beta)\) を更新していく。初期値は \((\alpha, \beta) = (0, 0)\) である。
- \(c_i\) が
A、\(c_{i-1}\) がAの場合：\(\alpha\) の値を \(\alpha + 1\) に更新する。- \(c_i\) が
A、\(c_{i-1}\) がBの場合：\(\alpha\) の値を \(\min(\alpha, \beta) + 1\) にした後、新しい \(\alpha\) の値を用いて \(\beta\) を \(\min(\beta, \alpha + 1)\) に更新する。- \(c_i\) が
B、\(c_{i-1}\) がAの場合：\(\beta\) の値を \(\min(\beta, \alpha) + 1\) に更新した後、新しい \(\beta\) の値を用いて \(\alpha\) を \(\min(\alpha, \beta + 1)\) に更新する。- \(c_i\) が
B、\(c_{i-1}\) がBの場合：\(\beta\) の値を \(\beta + 1\) に更新する。- ただし \(i = 1\) の場合は、\(c_{i-1}\) としてどちらを選んでも、次の \((\alpha, \beta)\) の値は同じになる。
最終的に、駅 \(N\) までの最短経路は \(\min(\alpha, \beta) + 1\) で求められる。
例えば \(c_i\) が A、\(c_{i-1}\) が B の場合は、下図のようなものが考えられます（丸の中の数字はここまでの最短経路長です）灰色の辺をそれぞれ順方向に通る場合、関係しない場合、逆方向に通る場合となっています。これでも、赤い層と青い層それぞれに限定して見たときは順方向とみなせるので、\(\alpha\) と \(\beta\) を順々に更新して答えを求めることができるのです。
ステップ 4：この問題の動的計画法による \(O(N^3)\) 解法
ステップ 3 では記録する状態数が \(N^2\) 通りになったので、あとはステップ 2 と同様にやるだけです！「駅 \(pos\) まで見たときに、状態が \((pos, \alpha, \beta, pre)\) になっているのは何通りか」を \(dp[pos][\alpha][\beta][pre]\) で記録します。（\(pre\) は直前の駅 \(pos\) を管理するグループが光速派なら 0、準急派なら 1 です）
\(pos = 0\) のときは、\(\alpha, \beta\) は駅 \(0\) を指しているので両方 \(0\) になります。また、最初 \(pre = 0\) としてよいので、最初の状態は \((pos, \alpha, \beta, pre) = (0, 0, 0, 0)\) とみなせます。アルゴリズム #4 と同様に、状態 \((pos, \alpha, \beta, pre)\) からは、次のように状態が変わります。
\(pre = 0\)（直前が光速派）のとき
- \(c_{pos+1}\) が
Aまたは?のとき、状態を \((pos+1, \alpha + 1, \beta, 0)\) に変えられる- \(c_{pos+1}\) が
Bまたは?のとき、状態を \((pos+1, \min(\alpha, \min(\beta, \alpha) + 1 + 1), \min(\beta, \alpha) + 1, 1) = (pos+1, \min(\alpha, \beta + 2), \min(\beta, \alpha) + 1, 1)\) に変えられる\(pre = 1\)（直前が準急派）のとき
- \(c_{pos+1}\) が
Aまたは?のとき、状態を \((pos + 1, \min(\alpha, \beta) + 1, \min(\beta, \min(\alpha, \beta) + 1 + 1), 0) = (pos + 1, \min(\alpha, \beta) + 1, \min(\beta, \alpha + 2), 0)\) に変えられる- \(c_{pos+1}\) が
Bまたは?のとき、状態を \((pos + 1, \alpha, \beta + 1, 1)\) に変えられる
このようにして \(dp[pos][\alpha][\beta][pre]\) の値を順々に求めていくことで、計算量 \(O(N^3)\) で本問題を解くことができます。最終的な答えは、\(\min(\alpha, \beta) \lt K\) であるようなすべての \((\alpha, \beta, pre)\) に対する \(dp[N-1][\alpha][\beta][pre]\) の合計になります。
サンプルコード (C++)： https://atcoder.jp/contests/arc119/submissions/22577361
サンプルコードでは \(pos\) の小さい順に \(dp[\alpha][\beta][pre]\) を更新していく実装をしています。なお、この計算量は \(O(N^3)\) なので、明らかに実行時間制限超過 (TLE) します。
ステップ 5：ひと工夫で計算量 \(O(N^2)\) に高速化
ステップ 4 で説明したアルゴリズムでは、DP の状態数が \(O(N^3)\) 通りになっています。計算量を減らすためには、この状態量を減らす必要があります。
ここで、下図の例を見てみましょう。下図では、\(\alpha = 4\) のまま \(\beta = 4, 5, 6, 7, 8, 9, \dots\) と増えていっています。だから、多くの種類の \((\alpha, \beta)\) を記録する必要があったのです。
しかし、次に光速派の駅が来ると、\(\beta\) はどれだけ大きくなっても \(6\) に更新されます。つまり、この例では、以下の図のように \(\beta\) が \(6\) 以上になったら \(\beta = 6\) としてみなしても良いのです。最終的には同じ最短経路長が得られます。
一般化すると、次のようになります。
- \(pre = 0\) のとき、\(\alpha \geq \beta + 2\) になったら、\(\alpha = \beta + 2\) とみなしても良い（すると、常に \(-2 \leq \beta - \alpha \leq +1\) が成り立つ）
- \(pre = 1\) のとき、\(\beta \geq \alpha + 2\) になったら、\(\beta = \alpha + 2\) とみなしても良い（すると、常に \(-1 \leq \beta - \alpha \leq +2\) が成り立つ）
そうすることで、限られた \((\alpha, \beta)\) に対してしか計算しなくても良くなります。全体で \(8N^2\) 通り程度の状態 \((pos, \alpha, \beta, pre)\) に対して DP の計算を行うことになります。このアイデアをステップ 4. の DP 解法に適用すると、以下のようになります。
\(pre = 0\)（直前が光速派）のとき
- \(c_{pos+1}\) が
Aまたは?のとき、状態を \((pos+1, \min(\alpha + 1, \beta + 2), \beta, 0)\) に変えられる- \(c_{pos+1}\) が
Bまたは?のとき、状態を \((pos+1, \min(\alpha, \beta + 2), \min(\beta, \alpha) + 1, 1)\) に変えられる\(pre = 1\)（直前が準急派）のとき
- \(c_{pos+1}\) が
Aまたは?のとき、状態を \((pos + 1, \min(\alpha, \beta) + 1, \min(\beta, \alpha + 2), 0)\) に変えられる- \(c_{pos+1}\) が
Bまたは?のとき、状態を \((pos + 1, \alpha, \min(\beta + 1, \alpha + 2), 1)\) に変えられる
実際の配列での計算では、メモリの制約の関係上 \(dp[pos][\alpha][\beta - \alpha + 2][pre]\) などの方法で値を管理することになるでしょう。このようにして、計算量 \(O(N^2)\) でこの問題を解けました！
6. 解法のまとめ
まとめると、この問題は以下の手順で解けます。
- \(dp[pos][\alpha][\beta][pre]\) を「駅 \(pos\) まで見たときに、直前の光速派・準急派の駅までの最短経路長がそれぞれ \(\alpha, \beta\)、直前の駅の状態が \(pre\)（0 = 光速派、1 = 準急派) である場合の数」として、これを順々に求めていくことを考える
- 逆方向の列車に乗る場合もあるが、それも普通列車に（2 連続せずに）乗るときに限定できるので、適切に \((\alpha, \beta)\) を更新していける
- そうすると、計算量 \(O(N^3)\) で答えが求まる
- 実際は \(\beta - \alpha \geq 2\) のとき \(\beta = \alpha + 2\) とみなせ、\(\alpha - \beta \geq 2\) のとき \(\alpha = \beta + 2\) とみなせるので、考えるべき \((\alpha, \beta)\) の状態数は \(O(N)\) 通りに減らせる
- そうすると、DP の計算量が \(O(N^2)\) に改善する
ここで、この解法を思いつくための重要な考え方は 2 つあります。
- 問題をより簡単にしたもの（番号の大きい方向にしか進めない場合）を考える
- 動的計画法による解法を考えるために、ベースの問題（最短経路を求める問題）をできるだけ少ない情報量 (= “メモリ使用量”) を記録することで解く方法を考える
1. の考え方はステップ 1, 2 で説明したものとなっていますが、元の問題よりも簡単に解法を思いつけるほか、問題の解法の手がかりにもなっています。
また、2. の考え方について、ステップ 3 では \(pos\) を含めた情報量が \(O(N^3)\) 通りになっていましたが、ステップ 5 で \(O(N^2)\) 通りに改善しています。動的計画法は、小さな部分問題に対する答えを順々に求めていくことで、全体の答えを求めるアルゴリズムなので、状態数を少なくすることが計算量改善の鍵になります。
この 2 種類の考え方は、この問題に限らず、他のプログラミングの問題を解くときにも有効な手段になりえます。
7. 実装例（サンプルコード）
この問題の解答例は以下の通りです。C++, Python による実装例を紹介します。
- Writer による実装 (C++)：https://atcoder.jp/contests/arc119/submissions/22307060
- Tester による実装 (Python)：https://atcoder.jp/contests/arc119/submissions/22356890
8. ＜発展＞ 理論的に高速な解法を求めて
実は、この問題は理論的には計算量 \(O(N \log^2 N)\) で解くことができます。発展的なトピックなので難しい言葉を用いて説明しますが、ご了承ください。
まず、ステップ 5 で説明した計算量 \(O(N^2)\) の解法は、ざっくりと言えば \(dp[pos][\alpha][state]\) を記録するもので、状態 \(state\) は \(8\) 種類の状態（それぞれ \((pre, \beta - \alpha) = (0, -2), (0, -1), (0, 0), (0, 1), (1, -1), (1, 0), (1, 1), (1, 2)\)）に対応させるものとします。
ここで、\(pos\) と \(state\) を固定して、DP の値を \(x\) の多項式 \(P_{pos, state}(x)\) として記録することを考えます。ここで：
\[P_{pos, state}(x) = \sum_{k=0}^{N-1} x^k \cdot dp[pos][k][state]\]
となるようにします。ここで DP 遷移の式を考えると、\(P_{pos}(x)\) の \(8\) 次元ベクトルから \(P_{pos+1}(x)\) の \(8\) 次元ベクトルへの変換は、\(8 \times 8\) 行列で表されます。ここで、行列の各成分は整数または \(ax+b\) の形になっています。
この変換に使われる行列は、\(c_i\) が A, B, ? である場合それぞれで違い、\(3\) パターンがあります。ここで、\(P_i(x)\) から \(P_{i+1}(x)\) に変換する行列を \(M_i\) とすると、最終的には次の値を求めることになります。
\[P_{N-1}(x) = M_{N-2}M_{N-3} \cdots M_2 M_1 M_0 \cdot P_0(x)\]
つまり、\(8 \times 8\) 行列 \(N-1\) 個の積を求めることになります。この積も \(8 \times 8\) 行列になり、各成分は最大で \(N-1\) 次式になります。なので、この行列積を普通に求めると、計算量 \(O(N^2)\) かかります。しかし、これは 分割統治法 と 高速フーリエ変換 (FFT) によって高速化することができます。
まず、分割統治法の使い方を説明します。分割統治法のアイデアは、半分ずつに分けて計算してこの結果をまとめるものです。例えば、\(N=17\) で、行列積 \(M_{15}M_{14} \cdots M_2 M_1 M_0\) を求めることを考えましょう。
- ステップ 1：\(M_1 M_0, M_3 M_2, M_5 M_4, \dots, M_{15} M_{14}\) の値を計算する
- ステップ 2：\(M_3 M_2 M_1 M_0, M_7 M_6 M_5 M_4, M_{11} M_{10} M_9 M_8, M_{15}M_{14}M_{13}M_{12}\) の値を、ステップ 1 の結果を使って計算する
- ステップ 3：\(M_7 M_6 \cdots M_2 M_1 M_0, M_{15}M_{14} \cdots M_{10} M_9 M_8\) の値を、ステップ 2 の結果を使って計算する
- ステップ 4：\(M_{15}M_{14} \cdots M_2 M_1 M_0\) の値を、ステップ 3 の結果を使って計算する
次に、高速フーリエ変換 (FFT) の使い方を説明します。FFT を使うと、2 つの \(N\) 次多項式 \(f(x), g(x)\) の積 \(f(x)g(x)\)を計算量 \(O(N \log N)\) で求められます。なお、FFT の詳しいアルゴリズムについては ATC 001-C：高速フーリエ変換 を見てください。
ここで、\(8 \times 8\) 行列 \(M_{r-1}M_{r-2} \cdots M_{l+2}M_{l+1}M_l\) の各成分は \(r-l\) 次以下の多項式になっています。ですから、ステップ \(k\) では \(2^{k-1}\) 次以下の多項式を成分とする行列の積を求めることになります。これに FFT を用いると、ステップ \(k\) 全体の計算量は \(O(2^k \cdot k) \times N/2^k = O(N \cdot k)\) となります。
よって、行列積を計算量 \(O(N \log^2 N)\) で解くことができました。すなわち、この問題を全体計算量 \(O(N \log^2 N)\) で解けたことになります。しかし、残念なことに定数倍が非常に大きいため、実用的には \(O(N^2)\) 解法よりも大幅に遅いと考えられます。
投稿日時:
最終更新: