【算法讲20：Dsu on Tree】树上数颜色 | Lomsat gelral

---

【例一：树上数颜色】

• 【链接】
  树上数颜色 | 洛谷 U41492
• 【题意】
  给定一棵有根树，有 $n$ 个节点，每个节点有一个颜色 $c_i$
  你需要求出每个节点，该节点的子树的颜色的种类数。
• 【范围】
  $1\le c_i\le n\le 10^5$

---

· 暴力做法 ·

• 对于每一个节点，遍历它的子树，计算出所有颜色的数量 $cnt[color]$，就能得到有多少个不同的颜色了。
• 时间复杂度：$O(n^2)$，结果 $TLE$

· 考虑优化 ·

• 我们可以算出每一个节点的重儿子。
  重儿子就是儿子中，子树大小最大的那一个儿子。其他的儿子都叫做轻儿子。
• 我们凭感觉，对于某一个节点，它的子树中的颜色数目的计算，可以先获取该节点的重儿子的子树的颜色数目，然后再去暴力跑其他轻儿子的子树的颜色数目，然后对这些数目进行合并。
  这样子是不是快的多呢？很神奇的，我们的时间复杂度从 $O(n^2)$ 变到了 $O(n\log n)$
• 略证：某个节点的遍历次数等于该节点到根节点的边中轻边的数量。从深度为 $d-1$ 的节点经过轻边走到深度为 $d$ 的节点，子树大小至少 $\times 2$，所以轻边个数为 $O(\log n)$，所以得证。

---

$\lceil DsuonTree\rfloor$

• 树上启发式合并，主要思想就是处理出每个节点的重儿子。接下来，我们保留重儿子的贡献，跑轻儿子的贡献，得到当前节点的贡献；然后，我们删掉轻儿子的贡献。
  主要是三个函数：$dfs1()$ 跑轻重儿子，$dfs2()$ 跑节点遍历，$go()$ 增/删贡献。
• 我们通过一些全局变量来模拟存储重儿子中的贡献值。
  然后好像就没了…

---

【代码】

• 时间复杂度：$O(n\log n)$
  空间复杂度：$O(n)$

vector<int>V[MAX];
int co[MAX];
int sz[MAX],hson[MAX];		/// 子树大小和重儿子编号。没有重儿子则编号为0
void dfs1(int x,int f){
    sz[x] = 1;
    for(auto it : V[x]){
        if(it == f)continue;
        dfs1(it,x);
        sz[x] += sz[it];
        if(sz[it] > sz[hson[x]])hson[x] = it;
    }
}
int res = 0;					/// 目前有多少种不同的颜色
int cnt[MAX],ans[MAX];			/// cnt[i] 表示颜色i的出现次数  ans[i] 表示结点编号i的答案
void go(int x,int f,int dec){
    cnt[co[x]] += dec;						/// 增  / 删贡献
    if(dec == 1 && cnt[co[x]] == 1)res++;	/// 多了一种颜色
    if(dec ==-1 && cnt[co[x]] == 0)res--;	/// 少了一种颜色
    for(auto it : V[x]){
        if(it == f)continue;
        go(it,x,dec);
    }
}
void dfs2(int x,int f,bool del){
    for(auto it : V[x]){
        if(it == f || it == hson[x])continue;	/// 遍历轻儿子
        dfs2(it,x,true);
    }
    if(hson[x])dfs2(hson[x],x,false);			/// 遍历重儿子
    cnt[co[x]]++;if(cnt[co[x]] == 1)res++;
    for(auto it : V[x]){
        if(it == f || it == hson[x])continue;	/// 增加轻儿子贡献
        go(it,x,1);
    }
    ans[x] = res;

    if(del)go(x,f,-1);		/// 删除轻儿子贡献
}
int main()
{
    int n;scanf("%d",&n);
    for(int i = 1;i < n;++i){
        int ta,tb;scanf("%d%d",&ta,&tb);
        V[ta].push_back(tb);
        V[tb].push_back(ta);
    }
    for(int i = 1;i <= n;++i)scanf("%d",&co[i]);
    dfs1(1,1);
    dfs2(1,1,false);
    int Q;scanf("%d",&Q);
    while(Q--){
        int t;scanf("%d",&t);
        printf("%d\n",ans[t]);
    }
    return 0;
}

---

【例二：$Lomsatgelral$】

• 【链接】
  洛谷 | CF 600E
• 【题意】
  给定 $n$ 个节点的树，每个节点的颜色为 $c_i$
  $an{s}_i$ 表示编号 $i$ 的子树中，所有出现次数最多的颜色的编号和。
  求出 $an{s}_1\sim an{s}_n$
• 【范围】
  $1\le c_i\le n\le 10^5$
• 【思路】
  和例题一差不多，仍然要计算出 $cnt[color]$
  不过我们还要存储 $mx$，表示目前子树的颜色的最多出现次数是多少
  然后去动态更新这个 $mx$，同时也要更新出 $res$ ，意义同 $an{s}_i$
• 【代码】
  时间复杂度：$O(n\log n)$
  空间复杂度：$O(n)$

vector<int>V[MAX];
int co[MAX];
int sz[MAX],hson[MAX];
void dfs1(int x,int f){
    sz[x] = 1;
    for(auto it : V[x]){
        if(it == f)continue;
        dfs1(it,x);
        sz[x] += sz[it];
        if(sz[it] > sz[hson[x]])hson[x] = it;
    }
}
ll mx = 0,res = 0;
ll cnt[MAX],ans[MAX];
void go(int x,int f,int dec){
    cnt[co[x]] += dec;
    if(dec == 1){		/// 增加贡献可能会增加 res
        if(cnt[co[x]] > mx)res = co[x],mx = cnt[co[x]];
        else if(cnt[co[x]] == mx)res += co[x];
    }
    for(auto it : V[x]){
        if(it == f)continue;
        go(it,x,dec);
    }
}
void dfs2(int x,int f,bool del){
    for(auto it : V[x]){
        if(it == f || it == hson[x])continue;
        dfs2(it,x,true);
    }
    if(hson[x])dfs2(hson[x],x,false);
    for(auto it : V[x]){
        if(it == f || it == hson[x])continue;
        go(it,x,1);
    }
    cnt[co[x]]++;
    if(cnt[co[x]] > mx)res = co[x],mx = cnt[co[x]];		/// 更新res mx
    else if(cnt[co[x]] == mx)res += co[x];
    ans[x] = res;

    if(del)go(x,f,-1),mx = res = 0;						/// 删除子树，mx = res = 0
}
int main()
{
    int n;scanf("%d",&n);
    for(int i = 1;i <= n;++i)scanf("%d",&co[i]);
    for(int i = 1;i < n;++i){
        int ta,tb;scanf("%d%d",&ta,&tb);
        V[ta].push_back(tb);
        V[tb].push_back(ta);
    }
    dfs1(1,1);
    dfs2(1,1,false);
    for(int i = 1;i <= n;++i)printf("%lld ",ans[i]);
    return 0;
}