Codeforces #1041 : E. Ancient Tree

ここまではすんなり解けたのだが。
https://codeforces.com/contest/2127/problem/E

問題

根付き木が与えられる。
また、各点はk色のいずれかで彩色するとする。

ある点vがcutieであるとは、vのsubtreeに以下の頂点対(x,y)が1個でもあることをいう。

• xとyの色は同じ
• xとvの色は異なる。
• xとyのLCAはvである。

各点のコストと、各点の色が与えられる。
ただし、一部の点は色が未確定であり、k色中任意の色を振れる。

cutieである頂点の総コストが最少となるような彩色方法を答えよ。

解法

葉から順に色を考えて行くと、Subtree内の色が増えると損をするので、未確定の点の色は、SubTree内に合わせるのが良い。
その際、Subtree内でcutieを満たす(x,y)の色が2通り以上あるときは、その点のコストがあるのは避けられない。
あとは、マージテクの要領でSubTree内の色の集合を管理しよう。

int T,N,K;
int W[202020];
set<int> S[202020];
int C[202020];
int NZ[202020];
vector<int> E[202020];
ll ret;

void dfs(int cur,int pre) {
	set<int> ng;
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur);
		if(S[e].size()>S[cur].size()) swap(S[cur],S[e]);
		FORR(a,S[e]) {
			if(S[cur].count(a)) ng.insert(a);
			S[cur].insert(a);
		}
	}
	if(ng.size()&&C[cur]==0) C[cur]=*ng.begin();
	if(ng.size()>=2 || (ng.size()==1&&C[cur]!=*ng.begin())) {
		ret+=W[cur];
	}
	if(C[cur]==0&&S[cur].size()) {
		C[cur]=*S[cur].begin();
	}
	if(C[cur]) S[cur].insert(C[cur]);
}

void dfs2(int cur,int pre,int c) {
	if(C[cur]==0) C[cur]=c;
	FORR(e,E[cur]) if(e!=pre) dfs2(e,cur,C[cur]);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N>>K;
		FOR(i,N) {
			cin>>W[i];
			E[i].clear();
			S[i].clear();
		}
		FOR(i,N) cin>>C[i];
		FOR(i,N-1) {
			cin>>x>>y;
			E[x-1].push_back(y-1);
			E[y-1].push_back(x-1);
		}
		ret=0;
		dfs(0,0);
		dfs2(0,0,1);
		cout<<ret<<endl;
		FOR(i,N) cout<<C[i]<<" ";
		cout<<endl;
		
		ll pret=ret;
		ret=0;
		FOR(i,N) S[i].clear();
		dfs(0,0);
		assert(pret==ret);
		
	}
}

まとめ

うーむ残念。