Codeforces #862 : Div2 E. There Should Be a Lot of Maximums

コードが意外にながいな。
https://codeforces.com/contest/1805/problem/E

問題

木を成す無向グラフを考える。
各点には整数値が設定されている。
木のMAD値とは、木で2回以上現れる整数値の最大値とする。

グラフから1辺取り除くと2つの木になるが、各辺について、その時の両木のMAD値の最大値を求めよ。

解法

元の木で3回以上登場する値は、必ずどちらかの木で2回以上現れるため、求めるMAD値の最大値はそれ以上である。
あとは木でちょうど2回現れる整数値が、辺の片側に現れるケースを考える。

これは、木をDFSしながら、SubTree内に1回または2回現れる整数値と、SubTree外に1回または2回現れる整数値を管理していけば、その最大値がわかる。

int N;
vector<int> E[101010];
int A[101010],D[2][101010];
int U[202020],V[202020];

map<pair<int,int>,int> M;
map<int,int> C;
int ret[202020];
multiset<int> Vs[202020];
vector<int> R;

void dfs(int cur,int pre,int id,int d) {
	D[id][cur]=d;
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur,id,d+1);
}
void dfs2(int cur,int pre,int id) {
	if(C[A[cur]]==2) Vs[id].insert(A[cur]);
	
	FORR(e,E[cur]) if(e!=pre) {
		if(D[0][e]+D[1][e]!=D[0][R[1]]) dfs2(e,cur,id);
	}
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>U[i]>>V[i];
		U[i]--,V[i]--;
		E[U[i]].push_back(V[i]);
		E[V[i]].push_back(U[i]);
		M[{U[i],V[i]}]=M[{V[i],U[i]}]=i;
	}
	int ma2=0,ma3=0;
	
	FOR(i,N) {
		cin>>A[i];
		C[A[i]]++;
	}
	FORR2(a,b,C) {
		if(b>=3) ma3=a;
		if(b==2) ma2=a;
	}
	FOR(i,N-1) ret[i]=ma3;
	
	if(ma2>ma3) {
		FOR(i,N) if(A[i]==ma2) R.push_back(i);
		dfs(R[0],R[0],0,0);
		dfs(R[1],R[1],1,0);
		vector<int> Rs(D[0][R[1]]+1);
		FOR(i,N) {
			if(D[0][i]+D[1][i]==D[0][R[1]]) {
				Rs[D[0][i]]=i;
				dfs2(i,i,i);
			}
		}
		FOR(i,N-1) {
			if(D[0][U[i]]+D[1][U[i]]!=D[0][R[1]]||D[0][V[i]]+D[1][V[i]]!=D[0][R[1]]) {
				ret[i]=ma2;
			}
		}
		set<int> L1,R1;
		map<int,int> L2,R2;
		FORR(r,Rs) FORR(v,Vs[r]) R2[v]=2;
		
		FOR(i,Rs.size()-1) {
			x=Rs[i];
			y=Rs[i+1];
			FORR(v,Vs[x]) {
				if(L1.count(v)) L2[v]=2;
				else L1.insert(v);
				if(R2.count(v)) {
					R2[v]--;
					if(R2[v]==1) {
						R2.erase(v);
					}
				}
			}
			k=M[{x,y}];
			if(L2.size()) ret[k]=max(ret[k],L2.rbegin()->first);
			if(R2.size()) ret[k]=max(ret[k],R2.rbegin()->first);
		}
	}
	
	FOR(i,N-1) cout<<ret[i]<<endl;
}

まとめ

今思えばもっと短く書けそう。