解きました。公式解説や有志の解説ブログを見ても自分と同じような考え方の人がいなかったので解説ブログを書きます。

問題リンク

https://atcoder.jp/contests/arc071/tasks/arc071_d

問題概要

$1,\ 2,\ 3,\ \dots,\ n$ からなる無限長の数列 $a_1,\ a_2,\ \dots$ のうち、次の条件を満たすものの個数を答えてください。

• $a_n=a_{n+1}=\dots$
• すべての $a_i$ について、$a_{i+1}=a_{i+2}=\dots=a_{i+a_i}$

制約

• $1\le n\le 10^6$

考察

見た目が DP っぽいので DP をします。$a_n$ を決めればその後も決まるため後ろから見たら良さそうです。まずは計算量を無視して、$O(N^2)$ でもいいので DP 出来ないか考えます。

素朴に考えると、$dp_{i,\ j}$ を $a_n,\ a_{n-1},\ \dots,\ a_{n-i+1}$ の $i$ 個まで見て、最後の $j$ 個先までが同じであるような数列の個数と定義したくなります。$a_n$ 以降も $1$ 個として数えます。しかし、この DP では上手く遷移を考えることが出来ません。理由は以下の通りです (本質ではなく長いので折りたたみました)。

よって、今まで見た数の先頭の数の次元 $k$ を追加して $dp_{i,\ j,\ k}$ とします ($k$ は先頭が $1$ なら $1$、そうでないなら $0$ です)。また、$a_n$ 以降も同じときは $j=\infty$ と扱います。このような DP を考えると上手に遷移できます。実際に試してみましょう。

さて、この DP の遷移図を先に示してしまいます。

$j\gt 1$ のときは $k=0$ となりえないことが分かりますね。
では、この遷移がそれぞれどのような意味であるのか見てみましょう。

• 初期条件
  $dp_{1,\ \infty,\ 0}=n-1,\ dp_{1,\ \infty,\ 1}=1$ です。

• $j=\infty$ からの遷移
  $dp_{i,\ \infty,\ k}=dp_{i+1,\ \infty,\ k}$ であることは容易に分かります。同じ数を追加するときです。
  そうでないとき、つまり違う数を追加するときは次の通りです。

  • $dp_{i+1,\ 1,\ 0}$ への遷移
    先頭の数が $1$ のときに別の数を追加するとき及び先頭の数が $1$ 以外のときに別の数を追加するときなので、$dp_{i+1,\ 1,\ 0}=(n-1)dp_{i,\ \infty,\ 0}+dp_{i,\ \infty,\ 1}$ です。
  • $dp_{i+1,\ 1,\ 1}$ への遷移
    先頭の数が $1$ 以外のときに $1$ を追加するときなので、$dp_{i+1,\ 1,\ 1}=dp_{i,\ \infty,\ 0}$ です。

  結局、$dp_{i,\ \infty,\ k}$ の値は一定なので、$dp_{i+1,\ 1,\ 0}$ には常に $(n-1)^2$ が、$dp_{i+1,\ 1,\ 1}$ には常に $n-1$ が足されます。

• $dp_{i,\ 1,\ 0}$ からの遷移
  この数列は先頭が $1$ 以外であり、その次は別の数なので、$1$ を追加するしかありません。よって $dp_{i+1,\ 1,\ 1}$ に遷移するほかありません。

• $dp_{i,\ j,\ 1}$ からの遷移
  この数列は先頭から $j$ 個 $1$ が続いているような状態です。よって、次の数は $1$ または $j$ 未満で $1$ でない数のどちらかです。よって、$dp_{i+1,\ j+1,\ 1}$ に等倍、$dp_{i+1,\ 1,\ 0}$ に $j-1$ 倍の遷移をします。

これで $O(N^2)$ の DP 解が求まりました。これを高速化してみましょう。

この遷移図をじっと見つめてみると、$j\gt 1$ で右下へ遷移するときは数が変わらないことに気付きます。よって、$dp_{n,\ j,\ k}$ をそれぞれ求めるときは $j\le 1$ で値が増えるときのみを考えれば良いです。
値が増えるのは、$j=\infty$ からの遷移と $j\gt 2$ からの遷移のふたつです。前者は簡単で、$dp_{n,\ 1,\ 0}$ のみ $(n-1)^2$ で他は $(n-1)n$ です。後者については DP の各値に $1,\ 2,\ 3,\ \dots$ を順番に掛けた値ですが、これは累積和の累積和のようにして $O(1)$ で求まります。

以上のようにして $dp_{n,\ j,\ k}$ の各値が求まるので、答えはその総和です。

解法

まず、以下の解法ではコーナーケースとなる $n=1,\ 2$ のときの答えはそれぞれ $1,\ 4$ です。以下では $n\ge 3$ とします。

考察とは違い、DP テーブルを $dp_i$ の一次元で定義し、これを $a_1,\ a_2,\ \dots,\ a_{n-i}$ の $n-i$ 個が同じ数である (が $a_n$ 以降は違う数である) ような場合の数とします。この値は考察で定義した DP では $dp_{n,\ n-i,\ 1}$ に相当します。
また、便宜上 $dp_n$ を $a_1\ne 1,\ a_1\ne a_2$ となるような場合の数とします。この値は考察で定義した DP では $dp_{n,\ 1,\ 0}$ に相当します。この値は $dp_{i,\ \infty,\ 0}$ からの遷移がないため、$n-1$ を引くというアドホックな補正が必要です。

まず、$dp_0=n-1,\ dp_1=dp_2=(n-1)n$ で初期化します。$i\ge 3$ となるときは次のように値を求められます。

• $\displaystyle dp_i=(n-1)n+\sum_{j=0}^{i-3} (i-2-j)dp_j$
  $=(n-1)n+(i-2)dp_0+(i-1)dp_1+\dots+dp_{i-3}$

さて、総和の部分ですが、$\sum_{j=0}^{i-3} dp_j$ 及び $\sum_{j=0}^{i-3} (i-2-j)dp_j$ を累積和の累積和のようにして更新すればこれは各 DP の値について $O(1)$ で計算可能です。

最後に、$dp_n$ から $n-1$ を引くことを忘れないでください。

以上の DP テーブルの総和と、全ての数が同じ場合の $n$ 通りを足せば答えです。

実装

		public void Solve()
		{
			var n = sr.ReadInt();
			if (n == 1)
			{
				Console.WriteLine(1);
				return;
			}
			if (n == 2)
			{
				Console.WriteLine(4);
				return;
			}
			var dp = new ModInt[n];
			ModInt ans = n;
			dp[0] = n - 1;
			dp[1] = (ModInt)(n - 1) * n;
			dp[2] = dp[1];
			ModInt sum = 0;
			ModInt sumall = 0;
			for (int i = 3; i < n; i++)
			{
				sum += dp[i - 3];
				sumall += sum;
				dp[i] += dp[1];
				dp[i] += sumall;
			}
			dp[n - 1] -= n - 1;
			for (int i = 0; i < n; i++)
			{
				ans += dp[i];
			}
			Console.WriteLine(ans);
		}

ACコード: https://atcoder.jp/contests/arc071/submissions/19515435

感想

二次元 DP を書いてそれを図を見ながら高速化する、という方針を取った人、他にいるんでしょうか…… (ほぼいない気がします)。しかし公式解説の方法とやっていることは恐らく同値ですし、非本質的な差だと思います。というか、公式解説のような普通の方法で考えた人は私の解説読んでもごちゃごちゃしてて分からない気がします。

まあでも、このような考え方をした人のさらなる理解の一助になれば幸いですので、今回解説を書いた次第です。相変わらず高難易度の問題は思考の経路を記しているので解説がごちゃごちゃしがちですが、後から私が見返して思い出せることを最優先にしています…… (当然ですが、そのような問題も簡潔に書けるようにもなりたいですね)。