[Codeforces600E]Lomsat gelral（dsu on the tree）

题目描述

传送门
题意：一棵树，每一个点有一个颜色，统计以每一个节点为根的子树中出现次数最多的颜色的编号和。

题解

学习了一下dsu on the tree，安利一下Yveh的良心博客：http://blog.csdn.net/qaq__qaq/article/details/53455462
算法的大致过程：
首先将树轻重链剖分，在dfs的过程中先dfs轻儿子，再dfs重儿子。
假设某一个点的儿子都已经被dfs过，统计这个点的答案。统计答案的过程中要calc当前这个点的子树，但是只calc它的轻链，重链不做。
这样的话，就需要在dfs的过程中，如果当前点是它父亲的轻儿子，做完这个点之后就将影响消除；而如果这个点是它父亲的重儿子，则将这个点的影响保留。
时间复杂度$O(nlogn)$，证明转自Yveh：
可以这么考虑：只有dfs到轻边时，才会将轻边的子树中合并到上一级的重链，树链剖分将一棵树分割成了不超过logn条重链。
每一个节点最多向上合并logn次，单次修改复杂度O(1)。
所以整体复杂度是O(nlogn)的。

代码

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
using namespace std;
#define N 100005
#define LL long long

int n,x,y,Son,Max;
int tot,point[N],nxt[N*2],v[N*2];
int c[N],size[N],son[N],cnt[N];
LL sum,ans[N];

void add(int x,int y)
{
    ++tot; nxt[tot]=point[x]; point[x]=tot; v[tot]=y;
}
void getson(int x,int fa)
{
    size[x]=1;
    for (int i=point[x];i;i=nxt[i])
        if (v[i]!=fa)
        {
            getson(v[i],x);
            size[x]+=size[v[i]];
            if (size[v[i]]>size[son[x]]) son[x]=v[i];
        }
}
void calc(int x,int fa,int val)
{
    cnt[c[x]]+=val;
    if (cnt[c[x]]>Max) sum=(LL)c[x],Max=cnt[c[x]];
    else if (cnt[c[x]]==Max) sum+=(LL)c[x];
    for (int i=point[x];i;i=nxt[i])
        if (v[i]!=fa&&v[i]!=Son) calc(v[i],x,val);
}
void dfs(int x,int fa,int k)
{
    for (int i=point[x];i;i=nxt[i])
        if (v[i]!=fa&&v[i]!=son[x]) dfs(v[i],x,0);
    if (son[x]) dfs(son[x],x,1),Son=son[x];
    calc(x,fa,1);Son=0;
    ans[x]=sum;
    if (!k) calc(x,fa,-1),Max=sum=0;
}
int main()
{
    scanf("%d",&n);
    for (int i=1;i<=n;++i) scanf("%d",&c[i]);
    for (int i=1;i<n;++i)
    {
        scanf("%d%d",&x,&y);
        add(x,y),add(y,x);
    }
    getson(1,0);
    dfs(1,0,0);
    for (int i=1;i<=n;++i) printf("%I64d%c",ans[i]," \n"[i==n]);
}