AtCoder ABC #302 (トヨタ自動車プログラミングコンテスト2023#2) : Ex - Ball Collector

ちょっと手間取った。
https://atcoder.jp/contests/abc302/tasks/abc302_h

問題

木を成す無向グラフが与えられる。
各点vには整数A[v]が書かれたボールとB[v]が書かれたボールの2つがある。

各点vに対し、以下を答えよ。

• 頂点1→vへのパス上の各点で、ボールを1つずつ選んだ時、選んだボールに現れる整数の種類数の最大値

解法

頂点v1個に対し、解を考える。
問題とは別のN頂点のグラフを考える。
パス上の点uに対し、A[u]-B[u]間に辺を張ったとする。

この時、各連結成分に対し

• 木である(辺の数=頂点数-1)ならば、(頂点数-1)種の整数を選べる
• 木でない(辺の数≧頂点数)ならば、(頂点数)種の整数を選べる

という取り方ができる。

よって連結成分毎の辺の個数と頂点数がわかれば、選べる整数の種類数がわかる。
あとは、根付き木をDFSしていき、巻き戻し可能なUnion-Findを使いならが、辺の個数や頂点数を更新していこう。

int N;
int A[202020],B[202020];
vector<int> E[202020];
int ret[202020];

template<int um> class UF_pop {
	public:
	vector<int> par,rank; vector<vector<int>> hist;
	UF_pop() {par=rank=vector<int>(um,0); for(int i=0;i<um;i++) par[i]=i;}
	void reinit() {int i; hist.clear(); FOR(i,um) rank[i]=0,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):operator[](par[x]);}
	void pop() {
		if(hist.size()) {
			auto a=hist.back();
			hist.pop_back();
			par[a[0]]=a[1]; rank[a[0]]=a[2];
			par[a[3]]=a[4]; rank[a[3]]=a[5];
		}
	}
	void operator()(int x,int y) {
		x=operator[](x); y=operator[](y);
		hist.push_back({x,par[x],rank[x],y,par[y],rank[y]});
		if(x==y) return;
		if(rank[x]<rank[y]) par[x]=y;
		else rank[x]+=(rank[x]==rank[y]), par[y]=x;
	}
};
UF_pop<202020> uf;
int NE[202020],NV[202020];
int num;
void dfs(int cur,int pre) {
	int a=uf[A[cur]];
	int b=uf[B[cur]];
	int pn=num;
	int pa=NE[a];
	int pb=NE[b];
	int va=NV[a];
	int vb=NV[b];
	if(a==b) {
		if(NV[a]-1==NE[a]) {
			num++;
			NE[a]++;
		}
	}
	else {
		if(NE[a]>=NV[a]&&NE[b]>=NV[b]) {
			;
		}
		else {
			num++;
		}
		NE[a]=NE[b]=NE[a]+NE[b]+1;
		NV[a]=NV[b]=NV[a]+NV[b];
		uf(a,b);
	}
	ret[cur]=num;
	FORR(e,E[cur]) if(e!=pre) dfs(e,cur);
	
	if(a!=b) uf.pop();
	num=pn;
	NE[a]=pa;
	NE[b]=pb;
	NV[a]=va;
	NV[b]=vb;
	
	
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	FOR(i,200001) NV[i]=1;
	cin>>N;
	FOR(i,N) {
		cin>>A[i]>>B[i];
		A[i]--,B[i]--;
	}
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	
	dfs(0,0);
	
	for(i=1;i<N;i++) cout<<ret[i]<<" ";
	cout<<endl;
	
}

まとめ

これもう少し早く思いつきたかったな。
二部マッチングをN回取ろうとして思いとどまってよかった。