Codeforces #599 Div1 D. Number Discovery

うーんややこしい。
https://codeforces.com/contest/1242/problem/D

問題

整数Kに対し、以下の手順で数列を作る。N番目の要素を求めよ。

• 数列中に含まれない正整数を小さい順にK個選び、順に数列に加える。
• 上記K要素の総和を数列に加え、また上に戻る。

解法

Editorialを見て解いたので簡単に。
正整数を、前者のステップで数列に加えるものと後者のすれっぷで数列に加えるものに分類する。
とすると(K^2+1)個ごとの区間に1つずつ後者のタイプが登場する。

Nが正整数を(K^2+1)個ごとの区間に分けた内の

• 後者タイプ
• 後者タイプの登場前
• 後者タイプの登場後

のいずれであるかを求められれば解も求まる。
各区間で後者のタイプとなる値は、再帰的に求めることができる。

int T;
ll N,K;

ll non_self(ll seg,ll K) {
	if(seg==0) return K*(K+1)/2;
	
	ll a=seg/K;
	ll b=seg%K;
	ll ns=non_self(a,K);
	ll s=a*(K*K+1)+b*K+1;
	ll t=a*(K*K+1)+b*K+K;
	ll add=0;
	if(t>=ns) add=min(K,t-ns+1);
	
	ll ans=a*K*(K*K+1)+b*K*K+K*(K+1)/2+add;
	return ans;
	
}

ll hoge(ll N,ll K) {
	ll seg=(N-1)/(K*K+1);
	ll ns=non_self(seg,K);
	
	if(ns==N) return (K+1)*(seg+1);
	
	if(N>ns) N--;
	N--;
	N%=(K*K+1);
	return seg*K*(K+1)+N/K*(K+1)+N%K+1;
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K;
		
		cout<<hoge(N,K)<<endl;
		
	}
}

まとめ

これ本番で詰められる気がしないな。