公式

F - Avoid Division 解説 by en_translator


In case of \(N=2\), i.e. if the tree has exactly two vertices, if some color \(i\) satisfies \(C_i\geq 2\) we may paint both vertices with that color; otherwise, no valid painting is possible.

Now assume \(N\geq 3\).

Let \(L\) be the number of leaves (vertices with degree \(1\)) in \(T\).
We first claim that no valid painting is possible when if the sum of \(C_i\), over all colors \(i\) such that \(C_i\geq 2\), is strictly less than \(L\). This is because, in that case at least one leaf must be painted with a color with \(C_i=1\), but cutting the only edge adjacent to that leaf splits it into the subtree solely consisting of that leaf and another subtree consisting of all other vertices, which never share the same color.

Therefore, it is necessary that the sum of \(C_i\) for all colors \(i\) with \(C_i\geq 2\) must be at least \(L\). In fact, there is always a valid painting when this is satisfied.

First, given the tree \(T\), we can take a vertex \(X\) that satisfies the following conditions:

  • \(X\) is not a leaf.
  • When \(X\) (and all adjacent edges) is removed from \(T\) to split it into subtrees, each of them contains at most \(\frac{L}{2}\) leaves. Here, note that a vertex is not counted as “leaves” if it newly becomes a leaf within its subtree after the removal.

This can be obtained by, for example, Depth-First Searching (DFS) from an arbitrary vertex, and picking the first vertex you encounter, in post-order, such that the subtree rooted at it contains at least \(\frac{L}{2}\) leaves. The latter condition is naturally satisfied by the condition, and we can also prove that this vertex is not a leaf when \(N\geq 3\).

Let us classify the leaves of \(T\) by which subtree it belongs to when \(X\) is removed. There are always at least two groups, and each group contains at most \(\frac{L}{2}\) leaves.
For this vertex, perform the following for \(i=1,2,\ldots, K\) in order:

  • If \(C_i=1\) do nothing; advance to next \(i\).

  • If all leaves are already painted, immediately terminate the procedure.

  • If there is exactly one unpainted vertex, paint that vertex and \(X\) with color \(i\), and immediately terminate the procedure.

  • If none of the above applies, do the following:

    1. Pick a leaf from any group containing the most unpainted vertices, and paint it with color \(i\).
    2. Pick a leaf from any group, except for the group picked in step 1., containing the most unpainted vertices, and paint it with color \(i\).
    3. Repeat the following operation \((C_i-2)\) times: pick a leaf from a group containing the most unpainted vertices, and paint it with color \(i\). In this step, we may choose a group already chosen in steps 1. or 2., or within step 3. If we run out of unpainted leaf during this repetition, immediately terminate the procedure.

During the procedure, by the end of processing each \(i\), the number of unpainted leaves \(L'\), and the number of unpainted leaves \(M'\) within a group with most unpainted leaves, satisfy \(M'\leq \frac{L'+1}{2}\); thus, note that if there are two or more unpainted leaves, there are at least two groups containing unpainted leaves, which guarantees us to perform the operation above.

Afterward, paint the remaining vertices with arbitrary unused colors (possibly with \(C_i=1\)). This is always possible because \(C_1+C_2+\cdots+C_K\geq N\).

Once the vertices are painted this way, every leaf satisfies the following condition:

There is another vertex with a different color than itself, among the leaves in the groups other than the leaf itself belongs to, or \(X\).

Now consider cutting \(T\) at an edge into two subtrees. The part that does not contain \(X\) contains one or more leaves of \(T\), which constitute a part or entire group of leaves fdefined above. On ther other hand, the part that does contain \(X\) include \(X\) and all leaves in the other groups. Thus, for a leaf \(Y\) of \(T\) contained in the part not containing \(X\), there exists a vertex painted in the same color as \(Y\) in the part containing \(X\). Hence, this painting always satisfies the condition in the problem statement.

Thus, we successfully constructed a painting that satisfies the condition in the problem statement above.
One can obtain vertex \(X\) in \(O(N)\), and determine the paintings in \(O(K+N\log N)\) by using a priority queue for example, which which are fast enough.

Therefore, the problem has been solved.

Sample code in 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;
}

投稿日時:
最終更新: