Codeforces #530 Div1 C. Construct a tree

これは方針は迷わなかったけど実装に手こずった。
https://codeforces.com/contest/1098/problem/C

問題

整数N,Sが与えられるので、以下を満たす根付き木を構成せよ。

• 頂点数はNで、各頂点1～Nのラベルが振られている。1番の頂点が根で、それ以外の頂点は、親頂点に自身より小さい番号を持つ。
• 各頂点のSubtreeのサイズの総和はSである。
• 条件を満たす木のうち、「子頂点の最大個数」が最小である。

解法

Subtreeのサイズの総和、というと一見扱いが難しいが、これは各頂点の深さの総和と一致する。
よって、「子頂点の最大個数」が小さいほど、木が深くなりSubtreeのサイズの総和は大きくなる。

そこで、「子頂点の最大個数」を小さい順に探索し、Subtreeのサイズの総和の範囲にSを含むものを求めよう。
「子頂点の最大個数」がCとする。
1～(N-C)を１列につなげた後、(N-C)番の頂点に(N-C-1)～N番の頂点を葉としてつなげると、Cに対し得られる「各頂点の深さの総和」の最大値となる。
この状態で総和がSを超える可能性があるので、葉頂点を浅い方向に1つずつ動かしていこう。

深さDの位置に頂点がf(D)個あるとき、深さD+1にはf(D)*C個までしか頂点を置けないことに留意すること。

ll N,S;
ll F[101010],G[101010];

ll num[101010],ma[101010];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>S;
	
	F[1]=G[1]=N*(N+1)/2;
	
	if(S==F[1]) {
		cout<<"Yes"<<endl;
		for(i=2;i<=N;i++) cout<<(i-1)<<" ";
		cout<<endl;
		return;
	}
	
	for(i=2;i<=N-1;i++) {
		int d=1;
		ll C=1;
		ll lef=N;
		
		while(lef>0) {
			ll num=min(C,lef);
			F[i]+=num*d;
			lef-=num;
			d++;
			C*=i;
		}
		
		G[i]=(N-i)*(N-i+1)/2+i*(N-i+1);
		if(F[i]<=S && S<=G[i]) break;
	}
	
	if(i==N) return _P("No\n");
	ll K=i;
	
	for(i=1;i<=N-K;i++) {
		num[i]=1, S-=i;
		if(i==1) ma[i]=1;
		else ma[i]=min(ma[i-1]*K,1LL<<40);
	}
	num[N-K+1]=K;
	S=K*(N-K+1)-S;
	x=N-K+1;
	
	while(S>0) {
		while(num[x]==0) x--;
		num[x]--;
		
		if(S<x && num[x-S]<ma[x-S]) {
			num[x-S]++;
			S=0;
		}
		else {
			for(i=1;i<=N;i++) if(num[i]<ma[i] && x-i<=S) {
				num[i]++;
				S-=(x-i);
				break;
			}
		}
	}
	
	cout<<"Yes"<<endl;
	int p=1;
	for(i=2;i<=N;i++) {
		
		FOR(j,num[i]) cout<<p+(j/K)<<" ";
		p+=num[i-1];
	}
	cout<<endl;
	
}

まとめ

もっと簡単に実装する方法あるのかな。