CSAcademy Round #29 : E. Root Change

解けはしたけどだいぶ苦戦。もっときれいに書けるのかな。
https://csacademy.com/contest/round-29/task/root-change/

問題

木を成すグラフがある。
各頂点iを根としたとき、ある辺とその先につながるSubTreeを切り落としても木の高さが変わらないような辺が何本あるか求めよ。

解法

逆に切ってしまうと木の高さが変わるような辺の数を数え上げよう。
DFSで葉から順に、頂点vのSubTree内の最も深い頂点の深さD(v)と、切ってしまうと木の高さが変わる辺の数E(v)を求めていこう。

頂点vについて、子頂点中でD(v')が単独最大である場合、v→v'に向かう辺を切っても木の高さが変わるのでE(v)=E(v')+1。
逆にD(v')が同着で複数最大の場合、どちらかのSubTreeを切っても反対側が残るのでE(v)=0となる。

あとはこの情報をもとに、全方位DPの要領で親方向の辺に対しD(v),E(v)を求めていこう。

int N;
vector<int> E[101010];
int D[101010];
int MD[101010],nc[101010],tot[101010];
int ret[101010];

void dfs(int cur,int pre,int d) {
	D[cur]=MD[cur]=d;
	FORR(e,E[cur]) if(e!=pre) {
		dfs(e,cur,d+1);
		if(MD[e]>MD[cur]) {
			MD[cur]=MD[e];
			nc[cur]=tot[cur]=0;
		}
		if(MD[e]==MD[cur]) {
			nc[cur]++;
			tot[cur]++;
			if(nc[e]==1) tot[cur]+=tot[e];
		}
	}
}
void dfs2(int cur,int pre,int md,int t) {
	
	vector<pair<int,int>> V;
	if(md>=0) V.push_back({md,t});
	
	FORR(e,E[cur]) if(e!=pre) {
		if(nc[e]==1) V.push_back({MD[e],1+tot[e]});
		else V.push_back({MD[e],1});
	}
	
	int mmd=max_element(ALL(V)).first;
	int tnc=0,ttot=0;
	FORR(r,V) if(r.first==mmd) {
		tnc++;
		ttot+=r.second;
	}
	
	if(tnc==1) ret[cur]=ttot;
	
	FORR(e,E[cur]) if(e!=pre) {
		if(MD[e]!=mmd) {
			if(tnc==1) dfs2(e,cur,mmd+2,1+ttot);
			else dfs2(e,cur,mmd+2,1);
		}
		else {
			if(tnc>2) {
				dfs2(e,cur,mmd+2,1);
			}
			else if(tnc==2) {
				dfs2(e,cur,mmd+2,1+ttot-(nc[e]==1?1+tot[e]:1));
			}
			else {
				int m2md=-1<<30,nc=0;
				int t=0;
				FORR(r,V) if(r.first!=mmd && r.first>m2md) m2md=r.first;
				FORR(r,V) if(r.first==m2md) t+=r.second, nc++;
				if(nc==1) dfs2(e,cur,m2md+2,1+t);
				else dfs2(e,cur,m2md+2,1);
			}
		}
	}
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,-1,0);
	dfs2(0,-1,0,0);
	FOR(i,N) cout<<N-1-ret[i]<<endl;
}

まとめ

全方位木DP苦手。