Official

F - Avoid Division Editorial by mechanicalpenciI


\(N=2\)、すなわち木が \(2\) 頂点からなる場合は、 \(C_i\geq 2\) であるような色 \(i\) が存在する場合はその色で \(2\) 頂点の両方を塗れば良く、そうで無い場合、条件をみたすように頂点を塗ることは不可能です。

以下では、\(N\geq 3\) のときについて考えます。

\(T\) の葉(次数が \(1\) の頂点)の数を \(L\) とします。
このとき、\(C_i\geq 2\) であるような 色 \(i\) についての \(C_i\) の総和が \(L\) 未満であるとき、条件をみたすように頂点に色を塗ることは不可能です。 なぜならば、葉のうち \(1\) つ以上は \(C_i=1\) であるような色で塗る必要がありますが、このとき、その葉と唯一繋がっている辺を切ると、その葉のみからなる部分木とそれ以外の全頂点からなる部分木に分かれ、これは共通部分を持たないからです。

よって、\(C_i\geq 2\) であるような 色 \(i\) についての \(C_i\) の総和は \(L\) 以上である必要がありますが、実はこのとき条件を満たすように頂点を塗ることがつねに可能です。

まず、\(T\) に対して、以下の条件をすべてみたす頂点 \(X\) が取れます。

  • \(X\) は葉でない。
  • \(X\)(および \(X\) につながる辺)をすべて削除した時、 \(T\) はいくつかの部分木に分かれるが、この部分木に含まれる \(T\) の葉の数はそれぞれ \(\frac{L}{2}\) 以下である。ここで、「 \(T\) の葉の数」には、\(X\) を削除したことで部分木の中で新たに葉になった頂点を含まないことに注意せよ。

これは例えば、適当な頂点を根として深さ優先探索 (DFS)等 を行い、その頂点を根とする部分木に属する葉の数が帰りがけ順で見て初めて \(\frac{N}{2}\) 以上となる頂点を選ぶことによって取得できます。後者の条件は頂点の取り方からみたしており、\(N\geq 3\) のときこれが葉とならないことも証明できます。

\(T\) の葉を、\(X\) を除いたときにどの部分木に属するかによってグループ分けします。グループは必ず \(2\) つ以上存在し、それぞれのグループに属する葉の数は高々 \(\frac{L}{2}\) です。
これに対して、\(i=1,2,\ldots, K\) の順で以下のことを行います。

  • \(C_i=1\) ならば何もせず、次の \(i\) へ進む。

  • すべての葉に色が塗られているならば操作をただちに終了する。

  • まだ塗られてない葉がちょうど \(1\) つのとき、その葉と \(X\) を色 \(i\) で塗り、操作をただちに終了する。

  • 上記のいずれでもないとき、次を行う。

    1. まだ塗られていない葉が最も多いグループの葉のうち \(1\) 枚を選び、色 \(i\) で塗る。
    2. 1. で選んだグループを除いたとき、まだ塗られていない葉が最も多いグループの葉のうち \(1\) 枚を選び、色 \(i\) で塗る。
    3. 「まだ塗られていない葉が最も多いグループの葉のうち \(1\) 枚を選び、色 \(i\) で塗る。」という操作を \(C_i-2\) 回繰り返す。ここでは、1., 2. で選んだグループや3. ですでに選んだグループを再度選んで良い。ただし、塗られていない葉が途中でなくなったならば、操作をただちに終了する。

このとき、操作の途中過程において、それぞれの \(i\) について終了した時点で、まだ塗られていない葉の総数を \(L'\)、まだ塗られていない葉が最も多いグループのその「まだ塗られていない葉」の数を \(M'\) とすると、\(M'\leq \frac{L'+1}{2}\) が成り立ち、特にまだ塗られていない葉が合計で \(2\) 枚以上存在するならば、まだ塗られていない葉が存在するグループが少なくとも \(2\) つ存在するため、上記の操作が必ず行えることに注意してください。

その後、残りの頂点を(\(C_i=1\) でもよい)、まだ使用可能な色を用いて適当に塗ります。これは \(C_1+C_2+\cdots+C_K\geq N\) よりつねに可能です。

さて、このように頂点に色を塗ったとき、任意の葉について以下が成り立ちます。

その葉の属するグループ以外のグループの葉あるいは \(X\) のうち、その葉と同じ色で塗られたものが存在する。

ここで、 \(T\) をある辺で切って \(2\) つの部分木に分けることを考えたとき、\(X\) を含まない方の部分木は \(1\) つ以上の \(T\) の葉を含みますが、これは先に分けたうちのある \(1\) つのグループの一部または全体となります。一方で、\(X\) を含む方の部分木は \(X\) および、そのグループ以外の全てのグループの葉をすべて含みます。よって、上の主張より、\(X\) を含まない方の部分木に含まれる\(T\) の葉の一つを \(Y\)とすると、\(Y\) と同じ色で塗られた頂点が \(X\) を含む方の部分木にも存在します。よって、このように頂点に色を塗ったとき、問題文の条件をつねにみたします。

よって、問題文の条件をみたす塗り方を構成することができました。
頂点 \(X\) の探索は \(O(N)\), 葉への塗り方はpriority_queue 等を用いることで \(O(K+N\log N)\) で行うことができ、十分高速です。

よって、この問題を解くことができました。

C++ による実装例:

#include <bits/stdc++.h>
using namespace std;

#define rep(i, n) for(int i = 0; i < n; ++i)
#define N (int)3e+5

int n,k;
vector<int>e[N];
int c[N];
int w[N];
int wsum,cent;
vector<int>llist[N];

void dfs(int p,int k){
	int sz=e[k].size();
	if(sz==1)w[k]=1;
	else w[k]=0;
	bool ok=true;
	rep(i,sz){
		if(e[k][i]!=p){
			dfs(k,e[k][i]);
			w[k]+=w[e[k][i]];
			int tmp=w[e[k][i]];
			if(sz==1)tmp++;
			if(tmp*2>wsum)ok=false;
		}
	}
	if(ok&&(wsum<=w[k]*2))cent=k;
	return;
}

void dfs2(int p,int k,int idx){
	int sz=e[k].size();
	if(sz==1)llist[idx].push_back(k);
	rep(i,sz){
		if(e[k][i]!=p){
			dfs2(k,e[k][i],idx);
		}
	}
	return;
}

int main(void){
	pair<int,int> pr,ps;
	int ans[N];
	int u,v;
	int t;
	cin>>t;
	rep(tt,t){
		cin>>n>>k;
		rep(i,n-1){
			cin>>u>>v;
			e[u-1].push_back(v-1);
			e[v-1].push_back(u-1);
		}
		rep(i,k)cin>>c[i];

		if(n==2){
			v=-1;
			rep(i,k)if(c[i]>=2)v=i+1;
			if(v<0)cout<<-1<<endl;
			else cout<<v<<" "<<v<<endl;
			rep(i,n){
				e[i].clear();
				llist[i].clear();
			}
			continue;
		}

		wsum=0;
		rep(i,n){
			if(e[i].size()==1)wsum++;
		}
		dfs(-1,0);
		int sz=e[cent].size();
		priority_queue<pair<int,int> >pq;
		rep(i,sz){
			dfs2(cent,e[cent][i],i);
			pq.push({llist[i].size(),i});
		}
		rep(i,n)ans[i]=-1;
		rep(i,k){
			if(wsum==0)break;
			if(c[i]>1){
				if(wsum==1){
					while(!pq.empty()){
						pr=pq.top();
						ans[llist[pr.second][0]]=i+1;
						c[i]--;
						wsum--;
						pq.pop();
					}
					ans[cent]=i+1;	
					c[i]--;
				}
				else if(wsum<=c[i]){
					while(!pq.empty()){
						pr=pq.top();
						rep(j,pr.first){
							ans[llist[pr.second][j]]=i+1;
							c[i]--;
							wsum--;
						}
						pq.pop();
					}
				}
				else{
					pr=pq.top();
					pq.pop();
					ps=pq.top();
					pq.pop();
					ans[llist[pr.second][pr.first-1]]=i+1;
					ans[llist[ps.second][ps.first-1]]=i+1;
					if(pr.first>1)pq.push({pr.first-1,pr.second});
					if(ps.first>1)pq.push({ps.first-1,ps.second});
					c[i]-=2;
					wsum-=2;
					rep(j,c[i]){
						pr=pq.top();
						pq.pop();
						ans[llist[pr.second][pr.first-1]]=i+1;
						if(pr.first>1)pq.push({pr.first-1,pr.second});
						wsum--;
					}
					c[i]=0;
				}
			}
		}
		if(wsum>0){
			cout<<-1<<endl;
		}
		else{
			int cur=0;
			rep(i,n){
				if(ans[i]<0){
					while(c[cur]<=0)cur++;
					ans[i]=cur+1;
					c[cur]--;
				}
				cout<<ans[i];
				if(i<(n-1))cout<<" ";
				else cout<<endl;
			}
		}
		rep(i,n){
			e[i].clear();
			llist[i].clear();
		}
	}
	
	return 0;
}

posted:
last update: