F: Two Snuke (AtCoder) - shimobepapaの日記

最近AtCoderという競技プログラミングをはじめたアラフォーサラリーマンです。ここのところ業務でコード書くことがなく寂しいので、ボケ防止のため在宅勤務で通勤がなくなってできた時間を使ってやっています。学生の人が多そうで肩身狭い思いしながらやっていますが、どこかで、始めた経緯とか感想などをまとめたいなと思います。

AtCoder、ABCという初心向けのはずのコンテストでも結構難しい問題多いです。解説も簡潔なものであることも多いので、解答読んでもすぐにはよくわからない・・ということがしばしばあります。同じように悩んでしている人いるかもしれないので、解答の理解に時間がかかった問題を丁寧に解説してみたいなと思います。（今週末コンテンストがないので暇潰しもかねて）

第一回は先週末にあったエイシングプログラミングコンテスト 2020の最後の問題です
F - Two Snuke
解説読んでも省略しているところで何をしているかわからず理解に苦労しました。

以下、解説のPDFの表記をそのまま使います。(Snukeの元ネタって誰なのかが気になります）
まず解説の「N-5個の一列に並んだボールに対し、 15個の仕切りを(最初の5つのセグメントに含まれるボールの個数が偶数になるように)入れる方法」を考えればいいというところへの誘導について考えていきます。

まず単純なケースとして
$\Delta s+ \Delta n+ \Delta u+ \Delta k+ \Delta e = N$
のとき
ありうる $(\Delta s, \Delta n, \Delta u, \Delta k, \Delta e)$ の組すべてについて
$\Delta s\cdot \Delta n\cdot \Delta u\cdot \Delta k\cdot \Delta e$
を計算して足し合わせた数について考えてみます。

これは「Nを5つのグループにわけ、それぞれのグループから１つ選ぶ」選び方の数に等しくなります。
５個の箱にそれぞれa, b, c, d, e個ボールが入っていて、それぞれの箱から１個ずつボールを選ぶボールの選び方は $a \cdot b \cdot c \cdot d \cdot e$ なので、総和を求める問題から組み合わせを求める問題に落とし込むことができています。（こんな発想はなかったので解説よんでおおーと思いました）

次に「N個のボールを５つのグループにわけ、それぞれのグループから１つボールを選ぶ」選び方を考えてみます。

f:id:shimobepapa:20200718012706p:plain
このボールを５つのグループにわけ、その中の一つを選んで色を塗ってみます
f:id:shimobepapa:20200718013103p:plain

ここで色を塗ったボールと仕切りに注目してみると、色塗りボール⇨仕切り⇨色塗りボール⇨仕切り・・・と交互に９個並んでいます。
そのため、色塗りボールを仕切りに変化させても元と１対１で対応します。（奇数番目を色塗りボールにすればいいので）

f:id:shimobepapa:20200718013558p:plain

ですので、「N-5個のボールに9個の仕切りを入れる、入れ方の個数」を求めればいいことがわかります。

いまはわかりやすくするため
$\Delta s+ \Delta n+ \Delta u+ \Delta k+ \Delta e = N$
という条件で考えましたが
$\Delta s+ \Delta n+ \Delta u+ \Delta k+ \Delta e \le N$
の場合はどうすればいいでしょうか？
この場合も同じ考え方で、
「N個のボールを６つのグループにわけ、最初の５つのグループからそれぞれ１つボールを選ぶ」
とすればいいです。
最後のグループがNに満たないあまりの数に対応します。
結果「N-5個のボールに10個の仕切りを入れる、入れ方の個数」となります。

さらに解説にあるように
$s + n + u + k + e + \Delta s+ \Delta n+ \Delta u+ \Delta k+ \Delta e \le N$
の場合も同じように考えれば、解説にあるように
「N-5個の一列に並んだボールに対し、 15個の仕切りを(最初の5つのセグメントに含まれるボールの個数が偶数になるように)入れる方法」
と１対１に対応しているところにたどり着きます。

で、自分が解説読んでつまづいたのが次からの部分です。
$dp(i,j,p)$ を「i個のボールまで見て、現在仕切りを j個入れ、現状のセグメントに含まれるボールの個数のパリティがであるような仕切りの入れ方の個数」、と定義してDPでこの問題を解くことができる。これを行列累乗で高速化する。

こちら最終的に理解した内容を書いてみます。

左から順番に、ボールか仕切りを置いていく。
と考えます。

最初５個のセグメント（仕切りが５つ置かれるまで）は偶数なので $j \le 5$ かそうでないかで場合分けして考えます。
解説にあるパリティというのは、いま置いて行っているセグメントにボールが偶数個入っているか、奇数個入っているかを示すもの、という意味だと思います。偶数の場合p=0、奇数の場合p=1と考えるとします。
なお簡単のため $j \ge 6$ のときは、奇数だったとしても常にp=0とします。
$j \le 5$ のとき
$dp(i, j, 0) = dp(i - 1, j, 1) + dp(i, j - 1, 0)\\ dp(i, j, 1) = dp(i - 1, j, 0)$

最後のセグメントのボールの数が偶数になるケース（p=0)は、奇数の状態にボールを置くか、仕切りを置くかです。ただ、かならず偶数にしなければいけないので、前のセグメントが奇数の状態で仕切りを置くことができません。
また奇数(p=1)になるケースは、偶数ボールがあるときにボールをもう１個おいたときのみです。

$j \ge 6$ のときは
$dp(i, j, 0) = dp(i - 1, j, 0) + dp(i, j - 1, 0)$
になります。ボールを置いた場合と、仕切りを置いた場合の場合の数の和になります

さて、これを行列の形にして解かなければいけません。
$v_n$ をベクトル、 $A$ を行列として
$v_n=Av_{n-1}$
という形にできれば
$v_n=A^nv_0$
となり、 $A^n$ は $o(log(n))$ で出せるため間に合います。

$dp(i, j, p)$ を少し定義をかえて、仕切りとボールをあわせてn個、仕切りをj個おいたときの数を $f(n, j, p)$ とします。
すると、
$j \le 5$ のとき
$f(n, j, 0) = f(n - 1, j, 1) + f(n - 1, j - 1, 0)\\ f(n, j, 1) = f(n - 1, j, 0)$
$j \ge 6$ のとき
$f(n, j, 0) = f(n - 1, j, 0) + f(n - 1, j - 1, 0)$

何か行列にできそうな気がしてきました。
ですので、考えられるj, pのパターンを全部いれたベクトルを考えることにします。
$v_n = \left(\begin{array}{c}f(n, 0, 0)\\ f(n, 0, 1)\\ f(n, 1, 0)\\ f(n, 1, 1)\\ .\\ .\\ .\\ f(n, 4, 0)\\ f(n, 4, 1)\\ f(n, 5, 0)\\ f(n, 6, 0)\\.\\ .\\ .\\ f(n, 15, 0)\end{array}\right)$

21要素ととても大きくなってしまいますが、プログラムで処理すれば実装も処理時間も全然余裕でした。

上の式を行列の形に書き直せば
$\left(\begin{array}{c}f(n, 0, 0)\\ f(n, 0, 1)\\ f(n, 1, 0)\\ f(n, 1, 1)\\ .\\ .\\ .\\ f(n, 4, 0)\\ f(n, 4, 1)\\ f(n, 5, 0)\\ f(n, 6, 0)\\.\\ .\\ .\\ f(n, 15, 0)\end{array}\right) = \left(\begin{array}{ccccccccccccccc}0 & 1 & 0 & 0 & . & . & . & 0 & 0 & 0 & 0 & . & . & . & 0\\1 & 0 & 0 & 0 & . & . & . & 0 & 0 & 0 & 0 & . & . & . & 0\\1 & 0 & 0 & 1 & . & . & . & 0 & 0 & 0 & 0 & . & . & . & 0\\0 & 0 & 1 & 0 & . & . & . & 0 & 0 & 0 & 0 & . & . & . & 0\\. & . & . & . & . & . & . & . & . & . & . & . & . & . & . \\. & . & . & . & . & . & . & . & . & . & . & . & . & . & . \\. & . & . & . & . & . & . & . & . & . & . & . & . & . & . \\ 0 & 0 & 0 & 0 & . & . & . & 0 & 1 & 0 & 0 & . & . & . & 0\\0 & 0 & 0 & 0 & . & . & . & 1 & 0 & 0 & 0 & . & . & . & 0\\ 0 & 0 & 0 & 0 & . & . & . & 1 & 0 & 1 & 0 & . & . & . & 0\\0 & 0 & 0 & 0 & . & . & . & 0 & 0 & 1 & 1 & . & . & . & 0\\. & . & . & . & . & . & . & . & . & . & 1 & 1 & . & . & . \\. & . & . & . & . & . & . & . & . & . & . & 1 & 1 & . & . \\. & . & . & . & . & . & . & . & . & . & . & . & 1 & 1 & . \\ 0 & 0 & 0 & 0 & . & . & . & 0 & 0 & 0 & 0 & . & . & 1 & 1\end{array}\right) \left(\begin{array}{c}f(n - 1, 0, 0)\\ f(n - 1, 0, 1)\\ f(n - 1, 1, 0)\\ f(n - 1, 1, 1)\\ .\\ .\\ .\\ f(n - 1, 4, 0)\\ f(n - 1, 4, 1)\\ f(n - 1, 5, 0)\\ f(n - 1, 6, 0)\\.\\ .\\ .\\ f(n - 1, 15, 0)\end{array}\right)$

となります。この行列を $A$ とします。
ボールも仕切りも何も置いてない状態は $n = 0, j = 0, p = 0$ なので
$v_0 = \left(\begin{array}{c}1\\ 0\\ 0\\ 0\\ .\\ .\\ .\\ 0\end{array}\right)$ となります。

求めたい答えは、ボールをN-5個、仕切りを15個並べた状態ですので、 $f(N+10, 15, 0)$ になります。 $v_{N+10}$ の最後の要素です。
ですので、
$v_{N+10} = A^{N+10}v_0$
を計算すれば答えがでてきます。

$A^n$ の高速計算の方法は色々なサイトで紹介していると思うので省略します。「繰り返し二乗法」で検索するのがいいと思います。

（軽い気持ちで書き始めましたが思ったより長くなってしまいました・・。第２回は果たしてあるのか）