kmjp's blog

競技プログラミング参加記です

AtCoder ARC #071 : F - Infinite Sequence

これ600-800pt位でもいいんじゃないかな…。
http://arc071.contest.atcoder.jp/tasks/arc071_d

問題

1~Nのいずれかの整数で構成される無限数列A[i]を考える。
以下の条件を満たすAは何通りあるか。

  • 第(N+1)項以降は第N項と一致する。
  • 第i項の後に続くA[i]項の値は一致する。

解法

以下A[i]は1-originとする。
以下の値を考える。
f(x) := 先頭N項のうち末尾x項、すなわちA[N-(x-1)]~A[N]の組み合わせ総数。

まず自明なものは以下の通り。

  •  f(x) = 1 (x≦0) : 第N項以降は第N項と同じ値を取る一択のため
  •  f(1) = N : A[N]は何をとってもよい
  •  f(2) = N^2 : A[N-1]、A[N]は何をとってもよい

以下、x≧3の時のf(x)を考えよう。yをN-(x-1)、すなわちN項めから数えて手前にx項目のとする。

  • A[y] = 1の時
    • この項の値の影響は、以後の項に影響しないのでf(x) += f(x-1)
  • A[y] = 2以上の時
    • 次の項A[y+1]に2以上の値を取った場合
      • 以後ずっと同じ値を取るしかなくなり数列の値が確定する。A[y]およびA[y+1]それぞれ2~Nの値を取る場合なのでf(x) += (N-1)^2
    • 次の項A[y+1]に1を取った場合
      • A[y+1]~A[y+A[y]]は1でなければならない。A[y+A[y]+1]以降はまた任意に設定できる。
      • よって \displaystyle f(x) += \sum_{i=2}^{N} f(x-(i+1))

f(x)を小さい順に計算し、その都度累積和を取っておけば、最後のsumもO(1)で計算でき、全体をO(N)で計算できる。

ll N;
ll mo=1000000007;
ll F[1010101];
ll S[1010101];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	
	F[1]=S[1]=N;
	F[2]=N*N%mo;
	S[2]=(S[1]+F[2])%mo;
	for(i=3;i<=N;i++) {
		F[i]=F[i-1];
		F[i]+=(N-1)*(N-1);
		F[i]+=S[i-3]+(N-1-(i-3));
		F[i]%=mo;
		S[i]=(S[i-1]+F[i])%mo;
	}
	
	cout<<F[N]<<endl;
	
}

まとめ

ARCの最終問題としてはだいぶ簡単目。