题目链接:https://www.luogu.com.cn/problem/CF600E
题意:一棵树有n个结点,每个结点都是一种颜色,每个颜色有一个编号,求树中每个子树的最多的颜色编号的和。
思路:我们首先处理出每个结点的重链,然后每次先处理轻链,计算贡献,然后再消去贡献,当最后处理到重链时
保留贡献,就这样启发式地合并,可以将复杂度降到nlogn。
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define fi first
#define se second
#define ls rt << 1
#define rs rt << 1|1
#define po pop_back
#define pb push_back
#define mk make_pair
#define lson l, mid, ls
#define rson mid + 1, r, rs
#define pll pair<ll, ll>
#define pii pair<int, int>
#define ull unsigned long long
#define pdd pair<double, double>
const int mod = 1e9 + 9;
const int maxn = 1e5 + 10;
const int inf = 0x3f3f3f3f;
const ll linf = 0x3f3f3f3f3f3f3f3f;
int c[maxn], sz[maxn], son[maxn], cnt[maxn], Max, Son;
vector<int> e[maxn];
ll sum = 0, ans[maxn];
void dfs(int u, int fa) //统计重儿子信息
{
sz[u] = 1;
for(auto v : e[u])
{
if(v == fa) continue;
dfs(v, u);
sz[u] += sz[v];
if(sz[son[u]] < sz[v]) son[u] = v;
}
}
void calc(int u, int fa, int val)
{
cnt[c[u]] += val;
if(cnt[c[u]] > Max) Max = cnt[c[u]], sum = c[u];
else if(cnt[c[u]] == Max) sum += c[u];
for(auto v : e[u])
{
if(v == fa || v == Son) continue;
calc(v, u, val);
}
}
void dfs2(int u, int fa, int op) // 0轻链 1重链
{
for(auto v : e[u])
{
if(v == fa || v == son[u]) continue;
dfs2(v, u, 0);
}
if(son[u]) //重儿子计算贡献 不消除
dfs2(son[u], u, 1), Son = son[u];
calc(u, fa, 1), Son = 0;
ans[u] = sum;
if(!op) calc(u, fa, -1), sum = Max = 0; //消除贡献
}
int main()
{
int n;
scanf("%d", &n);
for(int i = 1; i <= n; ++i) scanf("%d", &c[i]);
for(int i = 1; i < n; ++i)
{
int u, v;
scanf("%d%d", &u, &v);
e[u].pb(v), e[v].pb(u);
}
dfs(1, 0);
dfs2(1, 0, 0);
for(int i = 1; i <= n; ++i) printf("%lld ", ans[i]);
return 0;
}