HourRank19 : C. Maximal Tree Diameter

うーん、勿体ない。
https://www.hackerrank.com/contests/hourrank-19/challenges/maximal-tree-diameter

問題

木を成すグラフが与えられる。
ここから辺を1個取り除いて2つの木とし、さらに両方の木から1つずつ頂点を選んで辺を追加し、再び1つの木を作る。
こうして作る木の直径の最大値を求めよ。

解法

2つの木を作った場合、それらの直径を成す頂点の端の頂点同士を求めれば、最後(両方の木の直径の和+1)の直径を持つ木を作れる。
あとは、直径の和が最大となる2つの木を求めればよい。

これにはDFSを2回行うテクを用い、各頂点vに関し

• vのSubTreeにおける直径
• 元のグラフからvのSubTreeを除いたものにおける直径

を求めよう。
前者はSubTreeの最遠点と直径の最大値をそれぞれ更新する木DPを1回行えればよく、2回目で親方向の最大値を子に伝えていけばよい。

あとは根以外の頂点に対し、根に向かう辺を取り除いたと仮定すると、上記(2つの直径の和+1)が解の候補となる。

int N;
vector<int> E[505050];

int dia[505050],dep[505050];
vector<pair<int,int>> DI[505050];
vector<pair<int,int>> DEP[505050];

int ma=0;


pair<int,int> dfs(int cur,int pre) {
	FORR(r,E[cur]) if(r!=pre) {
		auto p=dfs(r,cur);
		p.second++;
		DI[cur].push_back({p.first,r});
		DEP[cur].push_back({p.second,r});
		dia[cur]=max(dia[cur],p.first);
		dia[cur]=max(dia[cur],dep[cur]+p.second);
		dep[cur]=max(dep[cur],p.second);
	}
	return {dia[cur],dep[cur]};
}

void dfs2(int cur,int pre,pair<int,int> par) {
	
	DI[cur].push_back({0,cur});
	DEP[cur].push_back({0,cur});
	DI[cur].push_back({par.first,pre});
	DEP[cur].push_back({par.second+1,pre});
	
	ma=max(ma, dia[cur]+1+par.first);
	
	sort(ALL(DI[cur]));
	sort(ALL(DEP[cur]));
	reverse(ALL(DI[cur]));
	reverse(ALL(DEP[cur]));
	
	FORR(r,E[cur]) if(r!=pre) {
		int di=0;
		vector<int> de;
		FORR(a,DEP[cur]) if(a.second!=r) {
			de.push_back(a.first);
			if(de.size()>=2) break;
		}
		FORR(a,DI[cur]) if(a.second!=r) {
			di=a.first;
			break;
		}
		if(de.size()==2) di=max(di,de[0]+de[1]);
		
		dfs2(r,cur,{di,de[0]});
	}
}


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);
	dfs2(0,-1,{-1<<20,-1<<20});
	cout<<ma<<endl;
}

まとめ

無駄に平方分割とかで頑張ろうとしてタイムロス。