AtCoder Regular Contest 071 F Infinite Sequence

arc071.contest.atcoder.jp

{1,…,n} からなる無限長の列 a₁,a₂,… のうち、 次の条件を満たしているものは何通りあるでしょうか？

• 第 n 項から先はすべて同じ数である。つまり、n≤i,j ならば a_i=a_j を満たす。
• どの正の整数 i に対しても、第 i 項の直後に並ぶ a_i 個の項はすべて同じ数である。つまり、 i<j<k≤i+a_i ならば a_j=a_k を満たす。

本番解けず。dp[k] = {k桁目から数字が連続する時の組み合わせ} を考える。

本番ではここでN桁目で残りk桁同じ場合、みたいなのを考えて解けなかった。

最初の2桁A Bに注目する。

• A=1のとき、B以降はdp[k - 1]通り。
• A≠1のとき、B≠1のとき、ABBBBB.... となる。(n-1)^2通り
• A≠1のとき、B≡1のとき、A111(A個1が並ぶ)111XXXXX となる。

最後の場合は、更に次の場合に分ける。

• A≧k - 1のとき A11111111... となる (k桁目以降は全部同じ数字) (N - k + 2)通り
• A≦k - 2のとき A1111(A個1が並ぶ) ときて、その後はdp[k-A-1] 通り。

最後のは要はdp[1]からdp[k - 3]の和になる。これをsum[k-3]とかくと次の漸化式がたち解ける。

dp[k] = dp[k - 1] + (N - 1)^2 + sum[k-3] + (N - k + 2)

sum[k]は累積和でO(N)、dpの計算もO(N)であり、答えはdp[N]となる。

ここで注意だが、dp[1]=Nはよいのだが、dp[2]の場合、ケース「A≧k - 1」でAが1も許される。そのため一番最初のケースとダブルカウントになるためsum[k-3]のところは-1に置き換える必要がある。これを気づかなくてeditorialの漸化式がなかなか合わなかった。