E. Ancient Tree
E. Ancient Tree
Bahamin came from the past to visit Ali — who came from the future. He also brought an ancient tree as a gift for Ali. He noticed some of its vertices have lost their color. Bahamin needs to repaint these vertices, but he is very busy with fixing his time machine. Fortunately (or unfortunately), dinosaurs now handle such tasks — for a fee, of course. He needs your help to find the coloring with minimum cost. So he gives you the problem as follows.
You are given a rooted tree$^{\text{∗}}$ of $n$ vertices, where vertex $1$ is the root. Each vertex has an integer weight $w_i$ and a color $c_i$, where the colors are integers between $1$ and $k$. However, some vertices have lost their colors, represented by $c_i = 0$.
We call vertex $v$ cutie if there exists two vertices $x$ and $y$, such that
- $\operatorname{lca}(x, y)$$^{\text{†}}$ $= v$,
- $c_x = c_y$, and
- $c_x \neq c_v$.
The cost of the tree is the sum of weights of all cutie vertices.
You have to assign colors between $1$ and $k$ to all the vertices which have lost their colors. Find the minimum possible cost among all valid colorings and provide a coloring with the minimum possible cost.
$^{\text{∗}}$A tree is a connected graph without cycles. A rooted tree is a tree where one vertex is special and called the root.
$^{\text{†}}$$\operatorname{lca}(x, y)$ denotes the lowest common ancestor (LCA) of $x$ and $y$.
Input
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 two integers $n$ and $k$ ($3 \leq n \leq 2 \cdot 10^5$, $2 \leq k \leq n$) — the number of vertices and the number of colors.
The second line contains $n$ integers $w_1,w_2,\ldots,w_n$ ($1 \leq w_i \leq 10^9$) — the weight of vertices.
The third line contains $n$ integers $c_1,c_2,\ldots,c_n$ ($0 \leq c_i \leq k$) — the color of vertices. $c_i=0$ means that vertex $i$ has lost its color.
Then $n−1$ lines follow, the $i$-th line containing two integers $u$ and $v$ ($1 \leq u, v \leq n$) — the two vertices that the $i$-th edge connects.
It is guaranteed that the given edges form a tree.
It is guaranteed that the sum of $n$ over all test cases does not exceed $2 \cdot 10^5$.
Output
For each test case, in the first line output a single integer — the minimum possible cost among all valid colorings.
In the second line output $n$ integers $c'_1, c'_2, \ldots, c'_n$ — a coloring with the minimum possible cost. You need to guarantee that:
- $c'_i = c_i$ if $c_i \neq 0$;
- $1 \leq c'_i \leq k$ if $c_i = 0$.
If there are multiple colorings with the minimum possible cost, you can output any of them.
Example
Input
4
4 4
5 5 5 5
1 0 2 3
1 2
1 3
1 4
5 2
3 1 4 1 5
1 2 1 2 2
1 4
2 1
3 4
4 5
11 3
3 1 4 3 1 4 3 1 4 5 6
0 0 0 2 1 2 1 2 2 1 1
1 2
2 3
2 4
2 5
2 6
1 7
7 8
7 9
10 3
3 11
4 3
2 3 2 3
2 1 0 0
3 1
1 2
2 4
Output
0
1 4 2 3
3
1 2 1 2 2
7
2 3 1 2 1 2 1 2 2 1 1
0
2 1 3 1
Note
In the first test case, there are four choices for the missing color:
- $c_2 = 1$ makes no vertex cutie, so the cost is $0$;
- $c_2 = 2$ makes vertex $1$ cutie, since $c_2 = c_3 = 2$, $\operatorname{lca}(2, 3) = 1$ and $c_1 \neq 2$. Thus, the cost is $w_1 = 5$;
- $c_2 = 3$ makes vertex $1$ cutie, since $c_2 = c_4 = 3$, $\operatorname{lca}(2, 4) = 1$ and $c_1 \neq 3$. Thus, the cost is $w_1 = 5$;
- $c_2 = 4$ makes no vertex cutie, so the cost is $0$.
Thus, the minimum possible cost among different colorings is $0$.
In the second test case, every vertex has a color. So current cost is not changeable. And since we have $c_5 = c_2 = 2$, $\operatorname{lca}(2, 5) = 1$ and $c_1 \neq 2$, vertex $1$ is cutie and current cost is equal to $w_1 = 3$.
In the third test case, a possible coloring with the minimum cost is shown below:

Some other valid colorings are:
- $c = [3, 1, 2, 2, 1, 2, 1, 2, 2, 1, 1]$ which makes vertices $1, 2, 3, 7$ cutie:
- $\operatorname{lca}(4, 8) = 1$;
- $\operatorname{lca}(3, 4) = 2$;
- $\operatorname{lca}(10, 11) = 3$;
- $\operatorname{lca}(8, 9) = 7$.
All pairs of vertices mentioned have a different color from their LCA. So the cost will be $w_1 + w_2 + w_3 + w_7 = 11$.
- $c = [3, 2, 1, 2, 1, 2, 1, 2, 2, 1, 1]$ which makes vertices $1, 2, 7$ cutie:
- $\operatorname{lca}(5, 7) = 1$;
- $\operatorname{lca}(5, 11) = 2$;
- $\operatorname{lca}(8, 9) = 7$.
All pairs of vertices mentioned have a different color from their LCA. So the cost will be $w_1 + w_2 + w_7 = 7$.
It can be shown that no coloring with cost smaller than $7$ exists.
解题思路
关键性质:对于某个未染色的点 $u$,如果其子树中存在染色的点,则选择子树中出现的颜色对 $u$ 染色,不会增加其祖先节点中 cutie 节点的数量。如果其子树中的点均未染色,则将整棵子树的点染成 $u$ 的父节点的颜色。下面给出证明。
先考虑子树中存在染色的点的情况:
- 如果在 $u$ 的两棵不同子树中分别存在点 $x,y$ 满足 $\text{lca}(x,y) = u$ 且 $c_x = c_y$,那么应该将 $u$ 染成 $c_x$。这是因为如果只有颜色为 $c_x$ 且满足 $\text{lca}$ 为 $u$ 的点对,此时将 $u$ 染成 $c_x$ 不会使得 $u$ 变成 cutie。而如果存在另外一个颜色不等于 $c_x$ 且满足 $\text{lca}$ 为 $u$ 的点对,那么不管将 $u$ 染成什么颜色,$u$ 都是 cutie。同时该染色方案不会导致 $u$ 的祖先变成 cutie。这是因为如果存在一个不在子树 $u$ 中的点 $v$,且其颜色 $c_v = c_x$,使得其祖先 $\text{lca}(u,v)$ 为 cutie,那么 $x$ 和 $v$ 同样会使得 $\text{lca}(x,v) = \text{lca}(u,v)$ 为 cutie。因此不管 $u$ 染成什么颜色 $\text{lca}(x,v)$ 都为 cutie。综上将 $u$ 染成 $c_x$ 不会使结果变糟糕。
- 如果在 $u$ 的子树中不存在点 $x,y$ 满足 $\text{lca}(x,y) = u$ 且 $c_x = c_y$,则让 $u$ 染成子树中任意出现过的颜色,假设为 $c_x$(其中 $x$ 是子树 $u$ 中的点且被染色)。首先 $u$ 本身不可能为 cutie,再考虑其祖先,同样地,如果存在一个不在子树 $u$ 中且满足 $c_v = c_x$ 的点 $v$ 使其祖先 $\text{lca}(u,v)$ 为 cutie,那么 $x$ 和 $v$ 也会使 $\text{lca}(u,v)$ 为 cutie。因此将 $u$ 染成 $c_x$ 不会使结果变糟糕。
再考虑子树中不存在染色的点的情况:此时应该将整棵子树 $u$ 的点都染成同种颜色,这样能保证整棵子树 $u$ 的点都不为 cutie。为了避免染色后使 $u$ 的祖先变成 cutie,应该将整棵子树染成 $u$ 的父节点 $\text{fa}(u)$ 的颜色。同样,如果存在一个不在子树 $u$ 中且满足 $c_v = c_{\text{fa}(u)}$ 的点 $v$ 使其祖先 $\text{lca}(u,v)$ 为 cutie,那么 $\text{fa}(u)$ 和 $v$ 也会使得 $\text{lca}(u,v)$ 为 cutie。因此将 $u$ 染成 $c_\text{fa}(u)$ 不会使结果变糟糕。
剩下的问题是怎么知道 $u$ 的子树出现了什么颜色?可以在 dfs 的过程中用 std::set 与启发式合并来维护子树中出现的颜色以及子树合并后的结果。在第一遍 dfs 中对子树中存在染色点的未染色点进行染色,并求出 cutie 节点的权值和。在第二遍 dfs 中对剩余未染色的点染成其父节点的颜色。剩下的细节请见代码。
AC代码如下,时间复杂度为 $O(n \log^2{n})$。
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 2e5 + 5, M = N * 2;
int a[N], c[N];
int h[N], e[M], ne[M], idx;
set<int> st[N];
LL ans;
void add(int u, int v) {
e[idx] = v, ne[idx] = h[u], h[u] = idx++;
}
void dfs1(int u, int p) {
st[u].clear();
set<int> vis;
for (int i = h[u]; i != -1; i = ne[i]) {
int v = e[i];
if (v == p) continue;
dfs1(v, u);
if (st[v].size() > st[u].size()) swap(st[u], st[v]);
for (auto &x : st[v]) {
if (!st[u].count(x)) st[u].insert(x);
else if (x != c[u]) vis.insert(x);
}
st[v].clear();
}
if (c[u]) {
st[u].insert(c[u]);
if (!vis.empty()) ans += a[u];
}
else if (!st[u].empty()) {
if (!vis.empty()) c[u] = *vis.begin();
else c[u] = *st[u].begin();
if (vis.size() > 1) ans += a[u];
}
}
void dfs2(int u, int p) {
if (!c[u]) c[u] = c[p];
for (int i = h[u]; i != -1; i = ne[i]) {
int v = e[i];
if (v == p) continue;
dfs2(v, u);
}
}
void solve() {
int n, m;
cin >> n >> m;
for (int i = 1; i <= n; i++) {
cin >> a[i];
}
for (int i = 1; i <= n; i++) {
cin >> c[i];
}
idx = 0;
memset(h, -1, n + 1 << 2);
for (int i = 0; i < n - 1; i++) {
int u, v;
cin >> u >> v;
add(u, v), add(v, u);
}
ans = 0;
dfs1(1, 0);
cout << ans << '\n';
if (!c[1]) fill(c + 1, c + n + 1, 1);
else dfs2(1, 0);
for (int i = 1; i <= n; i++) {
cout << c[i] << ' ';
}
cout << '\n';
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int t;
cin >> t;
while (t--) {
solve();
}
return 0;
}
参考资料
Atto Round 1 (Codeforces Round 1041, Div. 1 + Div. 2) A~E:https://zhuanlan.zhihu.com/p/1936961557159444775
Atto Round 1 (Codeforces Round 1041, Div. 1 + Div. 2) Editorial:https://codeforces.com/contest/2127/submission/333092546
本文来自博客园,作者:onlyblues,转载请注明原文链接:https://www.cnblogs.com/onlyblues/p/19030371

浙公网安备 33010602011771号