[Educational Codeforces Round 2 -E] Lomsat gelral

Codeforces传送门

洛谷传送门

题目大意

一棵以$1$为根的树有$n$个结点，每个结点都是一种颜色，每个颜色有一个编号$c_i(c\in [0,10^5])$，求树中以每个点作为根的子树中的最多的颜色编号的和。 （可能不止一种颜色，这时需要把这些颜色的编号加起来）

输入输出格式

输入格式

第一行一个正整数$n$。

第二行$n$个正整数， 表示每个节点的颜色编号。

以下$n-1$行， 表示树上的边。

输出格式

一行$n$个正整数， 第$i$个表示点$i$的子树中最多的颜色的编号和。

输入输出样例

输入样例#1：

4
1 2 3 4
1 2
2 3
2 4

输出样例#1：

10 9 3 4

输入样例#2：

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

输出样例#2：

6 5 4 3 2 3 3 1 1 3 2 2 1 2 3

解题分析

这类统计有根树子树内信息的题有个小trick，叫做$dsu\ on\ tree$（雾， 树上并查集？？？和并查集啥关系都没有啊）

实际上就是一种启发式合并。

这道题我们暴力做的话就是暴力统计每个子树的颜色个数， 或者从下向上递推， 时间复杂度为$O(N ^ 2)$， 但这样做显然做了很多重复的事。父节点的子树的就是由子节点的子树构成的， 所以我们可以想办法复用子节点的信息。

为了保证效率， 我们复用每个点重儿子的信息， 即始终将重儿子的信息保留在全局中。 具体流程如下：

• 递归DFS计算轻儿子及其子树中点的答案， 退出DFS时清空轻儿子的贡献；
• 递归DFS计算重儿子及其子树中点的答案， 退出时不清空；
• 计算轻儿子中的贡献， 合并入全局信息， 统计当前点的答案；

看起来似乎也很暴力， 但其实复杂度是$O(Nlog(N))\times$单点信息合并复杂度的。

因为我们用树剖剖出的轻重链都是$log(N)$条的， 而一个点的信息仅会在其处于一条轻链中的时候会被重复递归下去统计， 所以每个点统计到的次数也是$log(N)$次。

代码非常简单， 细节见下：

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cstdlib>
#include <cctype>
#include <cmath>
#define R register
#define IN inline
#define gc getchar()
#define W while
#define MX 100050
template <class T>
IN void in(T &x)
{
    x = 0; R char c = gc;
    W (!isdigit(c)) c = gc;
    W (isdigit(c))
    x = (x << 1) + (x << 3) + c - 48, c = gc;
}
struct INFO {int num; long long tot;} dat[MX];//存颜色出现个数及其标号总和
struct Edge {int to, nex;} edge[MX << 1];
int pos, dot, cnt, col[MX], son[MX], siz[MX], head[MX], cot[MX];
long long ans[MX];
bool hev[MX];//标记是否为重链
IN void addedge(R int from, R int to)
{edge[++cnt] = {to, head[from]}, head[from] = cnt;}
void DFS(R int now, R int fa)//找重儿子
{
    siz[now] = 1;
    for (R int i = head[now]; i; i = edge[i].nex)
    {
        if(edge[i].to == fa) continue;
        DFS(edge[i].to, now);
        siz[now] += siz[edge[i].to];
        if(siz[edge[i].to] > siz[son[now]]) son[now] = edge[i].to;
    }
}
void calc(R int now, R int fa, R int val)//计算贡献
{
    dat[cot[col[now]]].num--, dat[cot[col[now]]].tot -= col[now];
    cot[col[now]] += val;
    dat[cot[col[now]]].num++, dat[cot[col[now]]].tot += col[now];
    if(val > 0) pos = std::max(pos, cot[col[now]]);
    else if(!dat[pos].num) --pos;//答案是连续减少的， -1即可
    for (R int i = head[now]; i; i = edge[i].nex)
    {
        if(edge[i].to != fa && !hev[edge[i].to])
         //标记过的点即为统计中心的重儿子， 不需要计算了
        calc(edge[i].to, now, val);
    }
}
void DFS(R int now, R int fa, R bool typ)//递归处理
{
    for (R int i = head[now]; i; i = edge[i].nex)
    if(edge[i].to != fa && edge[i].to != son[now])
    DFS(edge[i].to, now, false);
    if(son[now]) DFS(son[now], now, true), hev[son[now]] = true;
    calc(now, fa, 1);
    ans[now] = dat[pos].tot, hev[son[now]] = false;
    if(!typ) calc(now, fa, -1);//清除贡献
}
int main(void)
{
    in(dot); int a, b;
    for (R int i = 1; i <= dot; ++i) in(col[i]);
    for (R int i = 1; i < dot; ++i)
    in(a), in(b), addedge(a, b), addedge(b, a);
    DFS(1, 0); DFS(1, 0, 1);
    for (R int i = 1; i <= dot; ++i) printf("%lld ", ans[i]);
}