yukicoder : No.1331 Moving Penguin

最初TLEが取れず手間取った。
https://yukicoder.me/problems/no/1331

問題

N個のマスが並んでおり、各マスには正整数A[i]が設定されている。
1番のマスにあるコマを以下のいずれかを繰り返しN番のマスに移動したい。
今x番のマスにコマがあるとき、

• (x+1)番のマスに移動する
• x ≡ y mod A[x]であるy＞xのマスに移動する。

移動の仕方は何通りあるか。

解法

f(n,k) := n番目のマスに至るとき、直前A[x]=kであるようなマスからyに移動してきたケースの組み合わせ
g(n) := n番目のマスに止まるケースの総数

とすると、以下の通り状態遷移できる。

• g(n) = sum(f(n,k)) + g(n-1)
• A[n]=1のとき
  • f(n+2,1) += g(n)
• それ以外の時
  • f(n+2,A[n]) += g(n)

ただ、f(n,k)を全kに対し行うとTLEするので平方分割を行う。
√N以上の大きいA[i]に対しては、その場でg(n+i*A[n])+=g(n)と結果を足しこんでしまおう。

int N;
int A[101010];
ll dp[101010][510];
ll S[101010];
const ll mo=1000000007;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	S[0]=1;
	for(i=1;i<=N;i++) {
		cin>>x;
		
		S[i]+=S[i-1];
		for(j=1;j<=500;j++) {
			S[i]+=dp[i][j];
			(dp[i+j][j]+=dp[i][j])%=mo;
		}
		S[i]%=mo;
		
		
		if(i==N) cout<<S[i]<<endl;
		
		if(x==1) {
			(dp[i+2][1]+=S[i])%=mo;
		}
		else if(x<=500) {
			(dp[i+x][x]+=S[i])%=mo;
		}
		else {
			for(j=i+x;j<=N;j+=x) S[j]+=S[i];
		}
		
	}
	
}

まとめ

これ想定解だったんだ。