Codeforces #699 Div2 : F. AB Tree

結構コード量増えちゃったな。
https://codeforces.com/contest/1481/problem/F

問題

N頂点の根付き木が与えられる。
このうち、X頂点に文字'a'、残りの(N-X)頂点に文字'b'を割り振ることを考える。

各頂点vに対し、文字列を以下のように構築する。

• 根からvまでのパス上にある頂点に割り振られた文字を、順に連結する。

構築できる文字列数が最小となるように文字の割り振り方を答えよ。

解法

基本的に、同じ深さにある頂点に同じ文字を割り振れば、構築する文字列はprefixが常に一致するため、文字列数は木の高さに一致する。
これはbitsetを使ったDPで、X個と(N-X)個の頂点への分割を容易に行える。

もしそのような(同じ深さにある頂点に同じ文字を割り振る)割り当てができない場合、どこかの深さの頂点で1か所'a''b'が共存するところを作るようにしよう。
また、その時はLeafである頂点群のみ'a''b'がばらける、逆に子頂点がある頂点群は'a'か'b'に統一しよう。

上記作業は貪欲に行える。
高さ順に、頂点群を'a'か'b'どちらかでそろえて行き、そろえられないときに上記処理を行う。

int N,K;
int leaf[101010];
vector<int> E[101010];
vector<int> C[101010];
vector<int> num[101010];

int mad;
bitset<101010> cur,add;
int tar[101010];

void dfs(int cur,int d) {
	mad=max(mad,d);
	C[d].push_back(cur);
	FORR(e,E[cur]) dfs(e,d+1);
	leaf[cur]=(E[cur].empty());
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	for(i=1;i<N;i++) {
		cin>>x;
		E[x-1].push_back(i);
	}
	dfs(0,0);
	
	MINUS(tar);
	cur[0]=1;
	FOR(i,N+1) if(C[i].size()) {
		add=cur<<C[i].size();
		bitset<101010> dif=add&(~cur);
		cur|=add;
		for(j=dif._Find_first();j<dif.size();j=dif._Find_next(j)) {
			tar[j]=i;
		}
	}
	string S(N,'b');
	if(cur[K]==1) {
		while(K) {
			x=tar[K];
			FORR(v,C[x]) S[v]='a', K--;
		}
		
		cout<<mad+1<<endl;
		cout<<S<<endl;
	}
	else {
		int X=K,Y=N-K;
		vector<int> lef;
		FOR(i,N+1) if(C[i].size()) {
			if(C[i].size()<=X) {
				FORR(v,C[i]) S[v]='a';
				X-=C[i].size();
			}
			else if(C[i].size()<=Y) {
				FORR(v,C[i]) S[v]='b';
				Y-=C[i].size();
			}
			else {
				vector<int> V;
				FORR(v,C[i]) {
					if(leaf[v]) lef.push_back(v);
					else V.push_back(v);
				}
				if(V.size()<=X) {
					FORR(v,V) S[v]='a';
					X-=V.size();
				}
				else if(V.size()<=Y) {
					FORR(v,V) S[v]='b';
					Y-=V.size();
				}
				else assert(0);
			}
		}
		FORR(v,lef) {
			if(X) S[v]='a', X--;
			else if(Y) S[v]='b', Y--;
			
		}
		
		cout<<mad+2<<endl;
		cout<<S<<endl;
	}
	
	
}

まとめ

この「考察すると実は貪欲で行けることがわかる」みたいな問題苦手だな…。