CodeTON Round 6 : F. Lazy Numbers

コードはそこまで長くないんだな。
https://codeforces.com/contest/1870/problem/F

問題

正整数N,Kが与えられる。
1～NをK進数表記し、それらを文字列とみなして昇順に並べる。
i番目にiが来ているようなiは何通りあるか。

解法

K進数表記のTrieを考える。
BFS順でx個めのノードbfs(x)にはxがある。
問題は、DFS順でx個めのノードdfs(x)にxがあるものをカウントする問題となる。

dfs(x)=bfs(x)となるxを、桁数毎に求める。
dfs(x)-bfs(x)は単調増加なので、dfs(x)-bfs(x)=0となる下限と上限を求めよう。

int T;
ll N,K;
__int128 L[61],R[61];
ll dfs(ll v) {
	ll num=-v+1;
	ll ov=v;
	//v以下の桁
	ll L=1;
	while(L<=v/K) L*=K;
	ll ol=L;
	while(v) {
		
		v/=K;
		L/=K;
		if(L) num+=v-L+1;
	}
	
	L=ol;
	ll R=ov-1;
	while(L<=R) {
		num+=R-L+1;
		if(L>N/K) break;
		L*=K;
		R=min(N,min(N/K+1,R+1)*K-1);
	}
	return num;
		
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K;
		
		ll ret=0;
		for(i=1;i<=61;i++) {
			L[i]=R[i-1]+1;
			R[i]=min((__int128)N,L[i]*K-1);
			if(L[i]>N) break;
			
			if(dfs(L[i])>0) continue;
			if(dfs(R[i])<0) continue;
			ll TL=L[i]-1;
			ll TR=L[i]-1;
			for(x=60;x>=0;x--) {
				if(TL+(1LL<<x)<=R[i]&&dfs(TL+(1LL<<x))<0) TL+=1LL<<x;
				if(TR+(1LL<<x)<=R[i]&&dfs(TR+(1LL<<x))<=0) TR+=1LL<<x;
			}
			if(TR>TL) ret+=TR-TL;
		}
		cout<<ret<<endl;
		
	}
}

まとめ

うーむ、これは思いつかなかった。