HackerRank University CodeSprint3 : E. Black White Tree

最後解けない問題を強引に定数倍高速化してTシャツ圏内。
https://www.hackerrank.com/contests/university-codesprint-3/challenges/black-white-tree

問題

木を成す無向グラフが与えられる。
各頂点は白か黒で塗られている。
このグラフの部分木のうち、白と黒の頂点数の差の絶対値が最大となるものを構成せよ。

解法

重心分解で解く。
重心を根とし、白が黒より多くなるケースと黒が白より多くなるケースをそれぞれDPで解いていく。
重心分解さえ気づけば基本的な木DPだね。

int N;
int C[101010];
vector<int> E[101010];
int vis[101010];
int NV[101010];
int dp[2][101010];

int dfs(int cur,int pre) {
	NV[cur]=1;
	FORR(e,E[cur]) if(e!=pre && vis[e]==0) NV[cur]+=dfs(e,cur);
	return NV[cur];
}

int dfs2(int cur,int pre,int TNV) {
	int tot=1;
	int ok=1;
	FORR(e,E[cur]) if(e!=pre && vis[e]==0) {
		int c = dfs2(e,cur,TNV);
		if(c!=-1) return c;
		tot += NV[e];
		if(NV[e]*2>TNV) ok=0;
	}
	if((TNV-tot)*2>TNV) ok=0;
	if(ok) return cur;
	return -1;
}

void dfs_count(int cur,int pre) {
	dp[0][cur]=dp[1][cur]=0;
	FORR(e,E[cur]) {
		if(e==pre) continue;
		if(vis[e]) continue;
		dfs_count(e,cur);
		dp[0][cur]+=max(0,dp[0][e]);
		dp[1][cur]+=max(0,dp[1][e]);
	}
	if(C[cur]==0) {
		dp[0][cur]=1+dp[0][cur];
		dp[1][cur]=max(0,dp[1][cur]-1);
	}
	else {
		dp[0][cur]=max(0,dp[0][cur]-1);
		dp[1][cur]=1+dp[1][cur];
	}
}

void split(int cur) {
	int TNV = dfs(cur,-1);
	if(TNV==0) return;
	int center=dfs2(cur,-1,TNV);
	
	dfs_count(center,-1);
	vis[center]=1;
	FORR(e,E[center]) if(vis[e]==0) split(e);
	
}

void dfs_set(int cur,int col,int pre,vector<int>& ret) {
	if(dp[col][cur]<=0) return;
	ret.push_back(cur+1);
	FORR(e,E[cur]) if(e!=pre) dfs_set(e,col,cur,ret);
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	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);
	}
	split(0);
	
	int id=-1,col=-1,ma=0;
	FOR(i,N) FOR(j,2) {
		if(dp[j][i]>ma) id=i, col=j, ma=dp[j][i];
	}
	
	ZERO(vis);
	vector<int> ret;
	dfs_count(id,-1);
	dfs_set(id,col,-1,ret);
	sort(ALL(ret));
	_P("%d\n",ma);
	_P("%d\n",ret.size());
	FOR(i,ret.size()) _P("%d%c",ret[i],(i==ret.size()-1)?'\n':' ');
	
	
}

まとめ

実装は若干面倒だけど、考察ではHackerRankの60ptにしては簡単目。