#P11194. County Capitals on a Tree
County Capitals on a Tree
County Capitals on a Tree
You are given a tree with (N) nodes (numbered from 1 to (N)). The nodes are partitioned into (K) counties. For each node (i), you are given an integer (B_i) ((1 \le B_i \le K)) indicating the county to which node (i) belongs. Your task is to choose one county capital for each county. Let the capitals be (A_1, A_2, \dots, A_K) (all distinct and each (A_j) is a vertex of the tree) such that the following condition holds:
[ \forall, 1\le i\le N, \quad \operatorname{dist}(i, A_{B_i}) < \min_{\substack{1\le j \le K \\ j \neq B_i}} \operatorname{dist}(i, A_j)]
Here, (\operatorname{dist}(a,b)) denotes the number of edges on the unique simple path between vertices (a) and (b). In other words, every node must be strictly closer to the capital of the county it belongs to than to any other county's capital.
Determine whether such an assignment of county capitals exists. If it exists, output one valid assignment; otherwise, output “No”.
inputFormat
The first line contains two integers (N) and (K) (the number of nodes and the number of counties).
The second line contains (N) integers (B_1, B_2, \dots, B_N) where (B_i) is the county number of vertex (i).
Then follow (N-1) lines, each containing two integers (u) and (v), indicating that there is an edge between vertices (u) and (v). It is guaranteed that the given graph is a tree.
outputFormat
If a valid assignment exists, print "Yes" in the first line. In the second line, print (K) integers (A_1, A_2, \dots, A_K), where (A_j) is the chosen capital for county (j). If no valid assignment exists, print "No".
sample
3 2
1 2 1
1 2
2 3
No