E - Zero-Sum Ranges 2 Editorial
by 
square1001
この問題は 900 点問題であり、本コンテストの中で最も難しい問題のうち 1 つです。
まず、\(+1\) と \(-1\) で構成される \(2^{2N}\) 通りの数列を全探索していると、制約が \(N \leqq 30\) なので実行時間制限には間に合いません。そのため本問題では、効率的な多項式時間アルゴリズムを構築することが求められます。
実際に、この問題は計算量 \(O(N^5)\) で解くことができます。以下のステップに分けて解説します。
- ステップ 1：「Zero-Sum Ranges」の問題をどう解くか
- ステップ 2：動的計画法のアイデア
- ステップ 3：簡単な場合を考える
- Zero-Sum Ranges 2 の \(O(N^5)\) 解法
- サンプルコード
1. 「Zero-Sum Ranges」の問題をどう解くか
まず、以下の問題を考えてみましょう。
長さ \(n\) の数列 \(a = (a_0, a_1, \dots, a_{n-1})\) が与えられる。
このとき、\(a_l + a_{l+1} + \cdots + a_{r-1} = 0\) となる \(l, r\) の組み合わせ \((1 \leqq l \leqq r \leqq N)\) は何個ありますか？
そうです、これは AtCoder Grand Contest 023 A「Zero-Sum Ranges」 という問題です。ここから、「Zero-Sum Ranges」の解法を少し説明します。
まず、全通りの \((l, r)\) についてそれぞれ判定していると、計算時間 \(O(n^3)\) 掛かってしまいます。しかし、これは累積和を使って以下のように工夫することができます。
数列 \(s_0, s_1, \dots, s_n\) を \(a_0, a_1, \dots, a_{n-1}\) の累積和とします。すなわち、\(i = 0, 1, \dots, n\) について、\(s_i = a_0 + a_1 + \cdots + a_{i-1}\) とします。このとき：
- \(a_l + a_{l+1} + \cdots + a_{r-1}\) は \(s_r - s_l\) で計算できる
つまり、区間の合計値が \(0\) となるのは、\(s_l = s_r\) となるとき、ということです。つまり、\(s_0, s_1, \dots, s_n\) から \(2\) つ選んだときに、同じ値となる組み合わせの個数が「Zero-Sum Ranges」の答えになっているのです。
上の図の例では、それぞれの段ごとに分けて考えて、合計が \(0\) となる区間は \(1 + 3 + 6 + 3 + 0 = 13\) 個、というように計算できる、といった感じです。
\(s_i = k\) となる \(i\) の個数が \(c\) 個の場合、これから \(2\) つを \(l, r\) として選ぶことで合計が \(0\) となる区間を作れるので、\(\frac{1}{2} c(c-1)\) 個の選び方があります。これを合計することで全体の個数が求められるという方法で、アルゴリズムとして高速です。また、詳しいことは次のステップで説明しますが、先ほどの図において「\(1\) 段ごとに見ることができる」という点が扱いやすく、この見方を本問題でも使うことができるのです。
2-1. 動的計画法のアイデア①
この問題を高速に解く際に、動的計画法を使うことはできないでしょうか。ここで、数列の先頭から見ていくことを考えましょう。
上の図は、数列の最初の \(7\) つが \(+1, +1, -1, +1, +1, +1, -1\) であることが決まっていて、次を決めようとしている状態を説明しています。現在、\((l, r) = (1, 3), (2, 4), (5, 7)\) で区間の合計が \(0\) になっている状態で、ここから \(1\) 個進むと以下のように分岐します。
- \(+1\) を置いた場合：\(s_8 = 4\) となり、\(l = 6\) について \(s_l = 4\) なので「合計が \(0\) の区間」は \(1\) 個増える
- \(-1\) を置いた場合：\(s_8 = 2\) となり、\(l = 2, 4\) について \(s_l = 2\) なので「合計が \(0\) の区間」は \(2\) 個増える
動的計画法は「現在の情報を（無駄を削って）まとめる」ことで計算量を削減していますが、どのような情報が必要なのでしょうか。先ほどの図の状態では、以下の情報を知っておかなければなりません。
- 数列をどこまで作ったか \(pos\)（図の場合 \(pos = 7\)）
- 現在までの「合計が \(0\) の区間」の数 \(cnt\)（図の場合 \(cnt = 3\)）
- 現在までの \(s_i = 4\) となる \(i\) の個数 \(c_4\)（図の場合 \(c_4 = 2\)）
- 現在までの \(s_i = 2\) となる \(i\) の個数 \(c_2\)（図の場合 \(c_2 = 1\)）
しかし、これだけの情報だけを記録して、その後に上手くいくはずがありません。\(s_i = 0, 1, 2, 3, 4, \dots\) の個数それぞれをすべて記録する必要があるのです。
すると、記録すべき状態数は \(N\) に対して指数的に増加します。確かに同じ \((pos, cnt, c_0, c_1, c_2, \dots)\) の組み合わせをまとめることができ、\(2^{2N}\) 通り全部を確かめる必要はないですが、\(N\) が増えるにしたがって状態数は \(5, 18, 56, 160, 432, 1120, 2816, \dots\) と増えていき、\(N = 17\) になれば \(12386034\) パターンと大変な量になってしまい、実行時間制限に間に合いません。
2-2. 動的計画法のアイデア②
先ほどは累積和の数列を左から順に作っていきましたが、今度は累積和の数列を図における「上から順に」作ってみることを考えてみましょう。
この例では、\(s_i = 2\) のところまで作っています。次のステップでは、⓪ の左、⓪ と ① の間、⑦ の右に「1」を挿入することになります。また、それ以外の場所に「1」を挿入することができないのも分かるでしょう。
現時点で、\(s_i = 4\) の位置が \(1\) 個、\(s_i = 3\) の位置が \(3\) 個、\(s_i = 2\) の位置が \(4\) 個あるので、今できている「合計が \(0\) の区間」は \(0 + 3 + 6 = 9\) 個です。また、上から数列を作っていくので、\(s_i = 4, 3, 2\) となる位置の個数がこれから増えることもありません。
この方法には、次のような重要な性質があるのです。
- 今までに作った段と、この後に追加される数が、相互に作用して「合計が \(0\) の区間」が新しく生まれることはない
つまり、「今までの \(s_i = 4\) の個数」などの情報を記録しなくても、新しく作られる「合計が \(0\) の区間」の数が分かるのです。1. 節の図で、段ごとに四角で囲った部分それぞれで「合計が \(0\) の区間の個数」が計算されているから、上の段から順に作っていくと今まで作った形の情報を考えなくて済むという点で都合がよい、というのが直感的な説明でしょう。
ここで少し整理してみましょう。上の図のような例で \(s_i = 2\) のところまで作り終わったとき、\(1\) を入れる必要があるのはどんな場所でしょうか？
- 数列の一番左（上図でいう「⓪ の左」の区間）
- 「\(s_i = 2\)」が \(2\) 連続している場所の間（上図でいう「⓪ と ① の間」の区間）
- 数列の一番右（上図でいう「⑦ の右」の区間）
これらの区間には最低 \(1\) 個は「\(1\)」を入れないと、全体が繋がらなくなってしまいますが、複数の「\(1\)」を入れられることにも注意してください。
ここで、同じ \(s_i\) が \(2\) 連続している部分を「穴」ということにしましょう。同じ \(s_i\) が \(3\) 連続している所では穴が \(2\) 個でき、そのどちらにも一個下の値をいくつか入れることになります。
3. 簡単な場合を考える
この節では、累積和 \(s_i\) がすべて \(0\) 以上の場合に限定して数え上げをすることを考えます。累積和が \(0\) 以上の場合と \(0\) 未満の場合で、以下のような違いがあります。
- \(s_i \geqq 0\) の時点では、「数列の一番左」「数列の一番右」に次の値を何個か追加しなければならない
- \(s_i < 0\) の時点では、「数列の一番左」「数列の一番右」に値を追加することはできない
ここで、\(s_i \geqq 0\) の時点で記録すべき必要な情報は、次の通りです。ただし、\(s_i \geqq level+1\) の場所を全部追加し終わってから、\(s_i = level\) となる場所を全部一気に追加するものとします。
- 現在の数列のサイズ \(size\)（上図の例だと \(size = 8\)）
- 最後に追加した \(s_i\) の値、これを \(level\) とする（上図の例だと \(level = 2\)）
- 現在作られた「合計が \(0\) の区間」の個数 \(cnt\)（上図の例だと \(cnt = 9\)）
- 現在の穴の個数 \(hole\)（上図の例だと \(hole = 4\)）
穴の具体的な位置関係などは記録する必要はなく、個数だけが今後の数列の運命を決める値となります。ここで、\(size \leqq 2N+1\)，\(level \leqq N\)，\(cnt \leqq N^2\)，\(hole \leqq N-1\) なので、状態数を \(O(N^5)\) 通りにすることができました。
さらに、よく考えてみると、\(level\) の値は記録する必要がありません。なぜなら、\(level\) の値がどうであれ、次への状態遷移は同一のものになるからです。よって、動的計画法では \(size, cnt, hole\) の値だけを記録することになり、状態数は \(O(N^4)\) 通りになります。
- \(dp[size][cnt][hole]\)：状態 \((size, cnt, hole)\) であるような数列が何通りあるか
数列を「上から順に」作っていくときの \((size, cnt, hole)\) の変化をにそって、\(dp[size][cnt][hole]\) の値を小さいものから順にどんどん計算していくことになります。
さて、ある状態 \((size, cnt, hole)\) から、一つ小さい数を一気に挿入するとき、どのような次の状態に行けるのか、考えてみましょう。例えば、下図の例は \((size, cnt, hole) = (8, 11, 3)\) の状態から一つ小さい数を挿入して、\((size, cnt, hole) = (17, 47, 4)\) の状態に遷移しています。
状態遷移では、上図におけるピンク色のエリアそれぞれに、一つ小さい数を \(1\) 個以上入れることになっています。ここで：
- ピンク色のエリアは \(hole + 2\) 個ある
- それぞれのエリアに \(1\) 個以上を入れることになる
ここで、ピンク色のエリアに \(x\) 個の数を入れる場合、次のステップでの穴の個数は \(x - (hole + 2)\) 個になります。これは次のように説明がつきます。
- 最初に各エリアに \(1\) 個ずつ、合計 \(hole + 2\) 個の数を入れると、穴の個数は \(0\) 個になる
- ここから余った数を \(1\) 個ずつ追加していくと、どのエリアに追加したかにかかわらず穴が \(1\) 個ずつ増えていく
また、ピンク色のエリアに \(x\) 個の数を入れる方法の数は、\(_{x-1}C_{hole+1}\) 通りになります。これは以下の問いの \(K = hole + 2, S = x\) の場合を考えれば分かります。
問い \(1\) 以上の整数の組 \((b_1, b_2, \dots, b_K)\) で、\(b_1 + b_2 + \cdots + b_K = S\) となるものは何通りありますか？
解答 \(b_1, b_1+b_2, b_1 + b_2 + b_3, \dots, b_1 + b_2 + \cdots + b_{K-1}\) には、\(1\) 以上 \(S-1\) 以下の異なる \(K-1\) 個の数が小さい順に登場することになります。\(1, 2, \dots, S-1\) から \(K-1\) 個を選ぶ方法は \(_{S-1}C_{K-1}\) 通りあるので、\((b_1, b_2, \dots, b_K)\) の組み合わせの数も \(_{S-1}C_{K-1}\) 通り、となります。
まとめると、\((size, cnt, hole)\) の状態から、次のステップで \(x\) 個 \((x \geqq hole + 2)\) 個の数を追加する場合：
- 状態が \((size + x, cnt + \frac{1}{2} x(x-1), x - (hole+2))\) に変わる
- このような状態遷移の方法は \(_{x-1}C_{hole+1}\) 通りある
ということが分かりました。これにしたがって動的計画法で \(dp[size][cnt][hole]\) を順々に求めていくと、計算量 \(O(N^5)\) ですべて求まります。
最終的な状態は \((size, cnt, hole) = (2N+1, K, 0)\) になっているので、「累積和 \(s_i\) がすべて \(0\) 以上」に限定した場合」の答えは \(dp[2N+1][K][0]\) になります。
4. ここからの解法
先ほどの考察で、\(s_i \geqq 0\) の場合に限定したときの答えを、計算量 \(O(N^5)\) で求めることができました。この方法を少し改良することで、\(s_i < 0\) を含めた場合にも解くことができるのか考えてみましょう。
\(1\) つ目の方針は、DP 配列の次元を \(1\) つ増やして、「現在 \(level \geqq 0\) か \(level < 0\) か」を 0/1 の値として記録することです。ここで、\(level < 0\) の場合、数を入れられるエリアは穴だけになる（両端には追加できなくなる）ので、次のステップで \(x\) 個 \((x \geqq hole)\) の数を追加する場合の状態遷移は以下のようになります。
- 状態が \((size + x, cnt + \frac{1}{2} x(x-1), x - hole)\) に変わる
- このような状態遷移の方法は \(_{x-1}C_{hole-1}\) 通りある
この点で、\(level \geqq 0\) の場合の状態遷移と \(level < 0\) の場合の状態遷移が異なることから、そのどちらの状態にあるのかを記録する必要があるのです。計算量は変わらず \(O(N^5)\) となります。
\(2\) つ目の方針は、\(level \geqq 0\) の場合と \(level \leqq -1\) の場合に分けて考えることです。まずは、下図の例を見てみましょう。
この例では：
- \(level \geqq 0\) の青い領域だけを見ると、\(level = 0\) 時点での \((size, cnt, hole) = (8, 11, 3)\) となっています。
- \(level < 0\) の赤い領域だけを見ると、\(level = -1\) 時点での \((size, cnt, hole) = (5, 6, 2)\) となっています。
ただし、\(level < 0\) に関しては、下から上の順番で作っていくものとします。そうすることで「\(level \geqq 0\) の場合と同様の \(dp[size][cnt][hole]\) の状態」として扱うことができます。
いま、上の部分（状態を \((size, cnt, hole)\) とする）と下の部分（状態を \((size', cnt', hole')\) とする）を組み合わせる際に、以下の条件が成り立っている必要があります。
- 合計の長さ \(size + size'\) が \(2N+1\) になっている
- 合計が \(0\) の区間の個数 \(cnt + cnt'\) が \(K\) になっている
- 両者の穴同士がうまくハマる必要がある、すなわち \(hole - hole' = 1\) になっている
よって、\((x, y, z)\) の組み合わせ全てに対する
- \(dp[x][y][z] \times dp[(2N+1)-x][K-y][z-1]\) を足し合わせたもの
が答えとなり、3. 節で説明した DP をそのまま使って計算することができます。全体計算量は変わらず \(O(N^5)\) になります。
5. 埋め込みについて
この問題の制約は \(1 \leqq N \leqq 30, 1 \leqq K \leqq N^2\) なので、入力としてありうる \((N, K)\) は \(1^2 + 2^2 + \cdots + 30^2 = 9455\) 種類しかありません。
また、答えは最大でも \(15\) 桁（\(N = 30, K = 230\) の場合の \(761128315856702\) が最大）なので、全部合わせても高々 \(9455 \times 15 = 141825\) 文字分にしかなりません。
ここで、本問題で提出できるプログラムの長さの上限は \(512\) KiB = \(524288\) バイトであり、必要な文字数よりも十分大きいです。したがって、ありうる入力すべてに対する答えの情報を、プログラムに “埋め込む” ことができるのです。
たとえあなたのプログラムの実行が少し遅かったり、プログラムの計算量が \(O(N^6)\) や \(O(N^7)\) であったとしても、多くの \((N, K)\) に対してまとめて答えを計算できるものであれば、手元のコンピューターであらかじめ全ての入力に対する答えを求めて、その出力結果を埋め込むことによって、実行時間制限超過 (TLE) を回避することができます。
しかしながら、これをやるためには全探索よりも十分に効率的なアルゴリズムを思いつく必要があり、その時点で正しい解法に近いものになっているでしょう。
6. サンプルコード
C++ における解答例は、次の通りです。
posted:
last update: