CF2252F Spectral Components

You are given a tree consisting of $ n $ vertices. Each vertex $ i $ is painted with a color $ c_i $ . For each distinct color $ c $ present in the tree, let $ m_c $ be the total number of vertices of color $ c $ . You are also given an array $ k $ of length $ n $ , where $ k_c $ ( $ 1 \le k_c \le m_c $ ) represents the target component size for color $ c $ . For every color $ c $ independently, your task is to choose a connected subgraph (a component) consisting of exactly $ k_c $ vertices. The vertices you choose for the component do not necessarily have to be of color $ c $ . The cost of a chosen component is the sum of the shortest distances from every vertex of color $ c $ to the chosen component. (The distance from a vertex $ v $ to a component $ S $ is defined as the minimum number of edges on a simple path from $ v $ to any vertex $ u $ in $ S $ ). For each color $ c $ from $ 1 $ to $ n $ , find the minimum possible cost of a valid component of size $ k_c $ . If there are no vertices of color $ c $ in the tree, output $ -1 $ for that color.

Each test contains multiple test cases. The first line contains the number of test cases $ t $ ( $ 1 \le t \le 10^4 $ ). The description of the test cases follows. The first line of each test case contains a single integer $ n $ ( $ 1 \le n \le 2 \cdot 10^5 $ ) — the number of vertices in the tree. The second line contains $ n $ integers $ c_1, c_2, \ldots, c_n $ ( $ 1 \le c_i \le n $ ) — the colors of the vertices. The third line contains $ n $ integers $ k_1, k_2, \ldots, k_n $ ( $ 1 \le k_i \le n $ ) — the target component sizes for each color. It is guaranteed that if color $ c $ appears $ m_c \gt 0 $ times in the tree, then $ 1 \le k_c \le m_c $ . Each of the next $ n - 1 $ lines contains two integers $ u $ and $ v $ ( $ 1 \le u, v \le n $ ), representing an edge between vertices $ u $ and $ v $ . It is guaranteed that the given edges form a valid tree. It is guaranteed that the sum of $ n $ over all test cases does not exceed $ 2 \cdot 10^5 $ .

For each test case, output $ n $ integers. The $ c $ -th integer should be the minimum possible cost of a valid component of size $ k_c $ for color $ c $ , or $ -1 $ if color $ c $ is not present in the tree.

In the first testcase, the tree has $ 5 $ vertices. Color $ 1 $ appears $ 3 $ times (vertices $ 1, 2, 4 $ ). Color $ 2 $ appears $ 2 $ times (vertices $ 3, 5 $ ). Colors $ 3 $ , $ 4 $ , and $ 5 $ do not appear, so their output is $ -1 $ . For color $ 1 $ ( $ k_1 = 2 $ ), we can choose the component $ S = \{2, 4\} $ . The distance from vertex $ 1 $ to $ S $ is $ 1 $ . The distances from vertices $ 2 $ and $ 4 $ to $ S $ are $ 0 $ . The total cost is $ 1 + 0 + 0 = 1 $ . For color $ 2 $ ( $ k_2 = 1 $ ), the optimal component is the single vertex $ S = \{2\} $ . The distance from $ 3 $ to $ 2 $ is $ 1 $ , and from $ 5 $ to $ 2 $ is $ 2 $ . The total cost is $ 3 $ . In the second testcase, the tree is a star graph with center $ 1 $ (color $ 2 $ ) and $ 5 $ leaves (color $ 1 $ ). For color $ 1 $ ( $ k_1 = 3 $ ), the optimal strategy is to include the center and two leaves, for instance, $ S = \{1, 2, 3\} $ . The distances from the color $ 1 $ leaves to $ S $ are $ 0 $ (for $ 2, 3 $ ) and $ 1 $ (for $ 4, 5, 6 $ ), yielding a minimum cost of $ 3 $ . For color $ 2 $ ( $ k_2 = 1 $ ), the only vertex is the center itself. Choosing $ S = \{1\} $ gives a cost of $ 0 $ . In the third testcase, the tree is a line graph $ 1-2-3-4-5-6 $ with alternating colors. For color $ 2 $ (vertices $ 2, 4, 6 $ ), we need a component of size $ 3 $ . The optimal component is $ S = \{3, 4, 5\} $ . The distances from the vertices of color $ 2 $ to $ S $ are $ 1 $ (from $ 2 $ , via edge $ 2-3 $ ), $ 0 $ (from $ 4 $ , since it is in $ S $ ), and $ 1 $ (from $ 6 $ , via edge $ 6-5 $ ). The total cost is $ 2 $ .