CSAcademy Round #20 : D. Stepping Number

手こずったけど最終的には解けきったのでよかった。
https://csacademy.com/contest/round-20/#task/stepping-number

問題

ある数がStepping numberであるとは、隣接する桁の数字の差の絶対値が1以下であるものを指す。
２整数N,Kが与えられる。
Nより大きなStepping numberのうち小さい順にK個目のものを求めよ。

解法

「Nより大きなStepping numberのうち小さい順にK個目のもの」を直接求めるのはしんどい。
そこでf(x) := x以下の非負整数のうち、Stepping numberである整数の数、という関数f(x)を考える。
そうすると、小さい順にf(N)+K番目のStepping numberを答えよという問題に置き換わる。

これらは事前に桁DPで
G(d,x) := 桁数がdで、最上位桁がxであるStepping numberの数
を事前に求めておけば、f(N)を求める処理もf(N)+K番目のStepping numberを求める処理も桁DPで上位桁から順に処理できる。

ll N,K;
ll dp[51][11];
string S;

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	dp[0][0]=dp[0][10]=1;
	FOR(i,10) dp[1][i]=1;
	dp[1][10]=10;
	for(i=2;i<=38;i++) {
		FOR(x,10) {
			dp[i][x]+=dp[i-1][x];
			if(x) dp[i][x]+=dp[i-1][x-1];
			if(x<9) dp[i][x]+=dp[i-1][x+1];
		}
		dp[i][10]=dp[i-1][10];
		for(x=1;x<=9;x++) dp[i][10]+=dp[i][x];
	}
	
	
	cin>>K>>S;
	
	if(S.size()==1) {
		K+=S[0]-'0'+1;
	}
	else {
		K += dp[S.size()-1][10];
		for(i=1;i<S[0]-'0';i++) {
			K += dp[S.size()][i];
		}
		for(i=1;i<S.size();i++) {
			if(S[i]>S[i-1]+1) {
				if(S[i-1]<'9') K += dp[S.size()-i][S[i-1]-'0'+1];
				K += dp[S.size()-i][S[i-1]-'0'];
				if(S[i-1]>='1') K += dp[S.size()-i][S[i-1]-'0'-1];
				break;
			}
			else if(S[i]<S[i-1]-1) {
				break;
			}
			
			if(S[i]==S[i-1]+1) {
				K += dp[S.size()-i][S[i-1]-'0'];
				if(S[i-1]>='1') K += dp[S.size()-i][S[i-1]-'0'-1];
			}
			if(S[i]==S[i-1]) {
				if(S[i-1]>='1') K += dp[S.size()-i][S[i-1]-'0'-1];
			}
		}
		if(i==S.size()) K++;
	}
	
	
	for(i=37;i>=0;i--) {
		if(K>dp[i][10]&&K<=dp[i+1][10]) {
			// i+1 digit
			int d=i+1;
			K -= dp[i][10];
			int S=1,E=9;
			FOR(x,d) {
				for(i=S;i<=E;i++) {
					ll sum=0;
					for(y=S;y<=i;y++) {
						sum+=dp[d-x][y];
					}
					if(sum>=K) {
						K-=sum-dp[d-x][i];
						cout<<i;
						S=max(i-1,0);
						E=min(i+1,9);
						break;
					}
				}
				
			}
			cout<<endl;
			return;
			
			
		}
	}
}

まとめ

方針はすぐに立ったのだが、デバッグにだいぶ苦労した。