読者です 読者をやめる 読者になる 読者になる

アルゴリズム忘備録

競技プログラミングとかデータ分析とか

AtCoder Regular Contest 071 F Infinite Sequence

arc071.contest.atcoder.jp

 

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

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

 

本番解けず。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の漸化式がなかなか合わなかった。