Codeforces #1023 : Div2 D. Apple Tree Traversing

なんかこの回割と苦労してるな。
https://codeforces.com/contest/2107/problem/D

問題

木を成す無向グラフが与えられる。
初期状態で各点には1個ずつリンゴが置いてある。

以下を繰り返したとき、得られる数列Aのうち辞書順最大のものを求めよ。

• パス上のリンゴが全部残っているような点u,vの距離をdとしたとき、Aに(d,u,v)を追加する。その際、パスu-v上のリンゴをすべて取り除く。

解法

直径のうち、(u,v)が辞書順最大となるものを選び、そのパスを取り除く、ということをリンゴがなくなるまで繰り返す。
(u,v)のうち辞書順最大なものを選ぶには、木の中心を求め、そこから最遠点のうち番号が多いもの上位2個を選んでいけばよい。

int T,N;
vector<int> E[202020];

vector<array<int,3>> ret;
int vis[202020];
vector<int> D[2];
vector<int> C;
int P[202020];

pair<int,int> farthest(int cur,int pre,int d,vector<int>& D) {
	C.push_back(cur);
	D[cur]=d;
	pair<int,int> r={d,cur};
	FORR(e,E[cur]) if(e!=pre&&vis[e]==0) r=max(r, farthest(e,cur,d+1,D));
	return r;
}

pair<int,int> dfs2(int cur,int pre,int d) {
	pair<int,int> D={d,cur};
	P[cur]=pre;
	FORR(e,E[cur]) if(e!=pre&&vis[e]==0) D=max(D,dfs2(e,cur,d+1));
	return D;
}

void diameter(int root) { // diameter,center
	if(vis[root]) return;
	C.clear();
	auto v1=farthest(root,root,0,D[0]);
	C.clear();
	auto v2=farthest(v1.second,v1.second,0,D[0]);
	C.clear();
	v1=farthest(v2.second,v2.second,0,D[1]);
	vector<int> R;
	int len=v1.first;
	FORR(a,C) if(D[0][a]+D[1][a]==len && abs(D[0][a]-D[1][a])<=1) R.push_back(a);
	if(len==0) {
		ret.push_back({1,R[0],R[0]});
		vis[R[0]]=1;
		return;
	}
	
	vector<pair<int,int>> V;
	if(R.size()==1) {
		int a=R[0];
		V.push_back({0,a});
		FORR(e,E[a]) if(vis[e]==0) V.push_back(dfs2(e,a,1));
		sort(ALL(V));
		reverse(ALL(V));
		int a2=V[0].second;
		int b2=V[1].second;
		ret.push_back({1+V[0].first+V[1].first,max(a2,b2),min(a2,b2)});
		while(a2!=a) {
			vis[a2]=1;
			a2=P[a2];
		}
		while(b2!=a) {
			vis[b2]=1;
			b2=P[b2];
		}
		vis[a]=1;
	}
	else {
		int a=R[0];
		int b=R[1];
		vis[b]=1;
		V.push_back(dfs2(a,a,0));
		vis[a]=1;
		vis[b]=0;
		V.push_back(dfs2(b,b,0));
		int a2=V[0].second;
		int b2=V[1].second;
		ret.push_back({2+V[0].first+V[1].first,max(a2,b2),min(a2,b2)});
		while(a2!=a) {
			vis[a2]=1;
			a2=P[a2];
		}
		while(b2!=b) {
			vis[b2]=1;
			b2=P[b2];
		}
		vis[a]=vis[b]=1;
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		FOR(i,N) {
			E[i].clear();
			vis[i]=0;
		}
		FOR(i,N-1) {
			cin>>x>>y;
			E[x-1].push_back(y-1);
			E[y-1].push_back(x-1);
		}
		ret.clear();
		D[0].clear();
		D[1].clear();
		D[0].resize(N);
		D[1].resize(N);
		FOR(i,N) while(vis[i]==0) diameter(i);
		
		
		
		
		sort(ALL(ret));
		reverse(ALL(ret));
		FORR(r,ret) cout<<r[0]<<" "<<r[1]+1<<" "<<r[2]+1<<" ";
		cout<<endl;
	}
}

まとめ

コードは長いけど、やることは自明なので難易度低め。