CF600E Lomsat gelral(树上启发式合并)
You are given a rooted tree with root in vertex 1. Each vertex is coloured in some colour.
Let's call colour c dominating in the subtree of vertex v if there are no other colours that appear in the subtree of vertex v more times than colour c. So it's possible that two or more colours will be dominating in the subtree of some vertex.
The subtree of vertex v is the vertex v and all other vertices that contains vertex v in each path to the root.
For each vertex v find the sum of all dominating colours in the subtree of vertex v.
Input
The first line contains integer n (1 ≤ n ≤ 105) — the number of vertices in the tree.
The second line contains n integers c**i (1 ≤ c**i ≤ n), c**i — the colour of the i-th vertex.
Each of the next n - 1 lines contains two integers x**j, y**j (1 ≤ x**j, y**j ≤ n) — the edge of the tree. The first vertex is the root of the tree.
Output
Print n integers — the sums of dominating colours for each vertex.
Examples
input
Copy
4
1 2 3 4
1 2
2 3
2 4
output
Copy
10 9 3 4
input
Copy
15
1 2 3 1 2 3 3 1 1 3 2 2 1 2 3
1 2
1 3
1 4
1 14
1 15
2 5
2 6
2 7
3 8
3 9
3 10
4 11
4 12
4 13
output
Copy
6 5 4 3 2 3 3 1 1 3 2 2 1 2 3
树上启发式合并典中典了属于是。dsu on tree整体的思路就是对于每个节点先搜一遍轻儿子且不保留统计信息,然后再搜重儿子且保留统计信息,之后再暴力搜一遍所有轻儿子进行统计,更新当前节点的答案。如果当前节点不是其父亲的重儿子的话则清除和当前子树有关的所有访问信息(包括访问当前节点的重儿子得到的信息)。把重儿子的信息留在最后的话就能为计算当前子树根节点的答案保留较多的信息量,这样就可以把\(O(n^2)\)的时间复杂度优化为\(O(nlogn)\)。实际写起来还是有很多细节需要注意的,详情见代码注释。
#include <iostream>
#include <cstring>
#define int long long
#define N 100005
using namespace std;
int n, c[N], head[N], ver[2 * N], Next[2 * N], tot = 0;
int sz[N], son[N], ans[N];
int cnt[N], mx = 0, sum = 0;
int hson = 0;//决定是否要 不遍历hson的标记(如果hson不为0,说明不要遍历标号为hson的节点)
void add(int x, int y) {
ver[++tot] = y, Next[tot] = head[x], head[x] = tot;
}
void dfs1(int x, int pre) {
sz[x] = 1;
int mxsz = -1;
for(int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y == pre) continue;
dfs1(y, x);
sz[x] += sz[y];
if(sz[y] > mxsz) {
son[x] = y;
mxsz = sz[y];
}
}
}
void process(int x, int pre, int v) {//暴力统计
cnt[c[x]] += v;//进行暴力统计的时候,cnt数组已经有访问重儿子得到的信息了
if(cnt[c[x]] > mx) {
mx = cnt[c[x]];
sum = c[x];
} else if(cnt[c[x]] == mx) {
sum += c[x];
}
for(int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y == pre || y == hson) continue;//注意hson是全局的,这样方便控制什么时候访问指定的重儿子 什么时候不访问
process(y, x, v);
}
}
void dfs2(int x, int pre, bool keep) {//keep为传递的重儿子标记
for(int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y == pre || son[x] == y) continue;
dfs2(y, x, 0);
}
if(son[x]) dfs2(son[x], x, 1), hson = son[x];
//此时重儿子对于sum的贡献(如果有的话)已经保留了
process(x, pre, 1);
//此时重儿子和轻儿子的信息都已经得到了 可以计算答案
ans[x] = sum;
hson = 0;//因为如果keep = 0的话需要消除得到的信息,hson这个这个重儿子的信息也需要消除,如果不重置的话process会避开,就无法完全消除了
if(!keep) {
//必须要先process再清空sum和mx,否则如果先清空sum,在process的时候sum还可能被更新,就无法达到清空的效果了
process(x, pre, -1);
sum = 0;
mx = 0;
}
}
signed main() {
cin >> n;
memset(son, 0, sizeof(son));
memset(sz, 0, sizeof(sz));
memset(cnt, 0, sizeof(cnt));
for(int i = 1; i <= n; i++) {
cin >> c[i];
}
for(int i = 1; i <= n - 1; i++) {
int x, y;
cin >> x >> y;
add(x, y);
add(y, x);
}
dfs1(1, 0);
dfs2(1, 0, 1);
for(int i = 1; i <= n; i++) {
cout << ans[i];
if(i != n) cout << " ";
}
return 0;
}