luogu p2664 树上游戏

lrb有一棵树，树的每个节点有个颜色。给一个长度为n的颜色序列，定义s(i,j) 为i 到j 的颜色数量。以及sum[i] = sigma(j = 1 to n) s(i,j)
现在他想让你求出所有的sum[i]

输入格式：
第一行为一个整数n，表示树节点的数量

第二行为n个整数，分别表示n个节点的颜色c[1],c[2]……c[n]

接下来n-1行，每行为两个整数x,y，表示x和y之间有一条边

输出格式：
输出n行，第i行为sum[i]

输入输出样例

输入样例#1：
5
1 2 3 2 3
1 2
2 3
2 4
1 5
输出样例#1：
10
9
11
9
12

这道题目首先可以联想到点分治，对于每一个分治树，我们首先求出重心，然后做一遍dfs求出重心的部分sum值，为什么说是部分值呢？因为答案可以分别有两种链的颜色贡献构成，第一种是分治树内的链的颜色贡献，第二种则是当前分治树外的链连接到当前点的颜色贡献。

我们可以设cnt[i]为颜色i对于重心的贡献，对于第一种，我们可以通过一遍dfs求出来，对于搜索到的每一个颜色，如果是第一次出现，我们则把cnt[i]加上这个点和其子树的大小。对于第二种贡献，我们可以通过当前点分治的上一次递归求出，具体地，对于上一次点分治的递归同样有cnt数组，性质同上，同时我们也求出了sum[重心]的答案（先不管这个重心的sum的第二种答案贡献是在哪里求解出来的），那么对于每一次划分更小的点分树之前，我们先用cnt数组来把某个子树的第二种贡献值求出来，也就是这个子树之外其他子树的链连接到这个子树的值。更具体的，对于每一个处理的子树，我们首先在cnt数组中将这个子树的贡献去除，得到的剩下的cnt数组就是其他子树对于重心的贡献，那么其他子树一定可以通过中心来连接到这个子树里的点，我们可以将除去这个子树贡献之后所剩下的总贡献作为一个基础值pr，即：该子树内的任意一个点的第二种答案贡献都不会小于这个基础值pr。
然后我们就可以对这个子树进行dfs了，具体地，对于每一个点x，我们首先把这个点x的sum值加上基础值pr ：sum[x] += pr+addtion（一个dfs的递归值，初始为0）
然后再进行判断，如果这个点x的颜色是子树根到该点的路径上第一次出现，那么点x的sum会额外加k（ k = 整个分治树大小与该子树大小的差值再减去cnt数组中x的颜色的贡献），可以理解为点x的对于该子树外的所有路径（链）中，有些点是包含了x的颜色的，有些点是不包含的，那么k就是剩下不包含x的颜色的路径（链），加上k就使得点x对于该子树外的所有点的路径（链）都包含了点x的颜色的贡献在里面。然后在下一步dfs递归的过程中，递归参数addtion要加上k。至此，第二种答案贡献值就求解出来了，值得注意的是，第二种答案不是在当前分治树中求解出来，而是在这之前的上一个递归的分治树中已经求解出来了。

看着挺长的。。。有点不容易理解，这个模拟一遍大概就好了吧，233。点分的板子可以略过，主要是cnt数组的作用和第二种答案贡献值是在上一个递归中求解出来的以及第二种贡献的求解方法步骤理解了就好写了。。

下面贴代码：

#include <bits/stdc++.h>
#define ll long long int
using namespace std;
vector <int> t[110000];
int n,wh,mi,sz,to,lt[110000],ext[110000];
int sign[110000],tot[110000];
ll cnt[110000],sum[110000],cnt2[110000];
int col[110000],size[110000],ape[110000];

int wt(int s,int siz) {
    if(abs((size[s]-1)-(siz-size[s])) < mi) {
        wh = s; mi = abs((size[s]-1)-(siz-size[s]));
    }
    sign[s] = 1;
    int k = t[s].size();
    for(int i = 0;i < k;i++) {
        if(sign[t[s][i]] == 0) {
            wt(t[s][i],siz);
        }
    }
    sign[s] = 0;
    return 0;
}

int up_size (int s) {
    sign[s] = 1;
    int st = 1;
    int k = t[s].size();
    for(int i = 0;i < k;i++) {
        if(sign[t[s][i]] == 0) {
            st += up_size(t[s][i]);
        }
    }
    size[s] = st;
    sign[s] = 0;
    return st;
}

int dfscore(int s,int fi, int nc) {
    if(ext[col[s]] == 0) {
        ext[col[s]] = 1;
        sz++; lt[sz] = col[s];
    }
    sign[s] = 1;
    ape[col[s]]++;
    if(ape[col[s]] == 1) {
        cnt[col[s]] += size[s];
        sum[fi] += size[s];
        nc++;
    }
    tot[s] = nc;
    int k = t[s].size();
    for(int i = 0;i < k;i++) {
        if(sign[t[s][i]] == 0) {
            dfscore(t[s][i],fi,nc);
        }
    }
    ape[col[s]]--;
    sign[s] = 0;
    return 0;
}

int dfs1(int s,int np,int fi) {
    sign[s] = 1;
    ape[col[s]]++;
    if(ape[col[s]] == 1) {
        if(np == 1) {
            cnt[col[s]] += size[s];
            to-=size[s];
        } else {
            cnt[col[s]] -= size[s];
            to+=size[s];
        }
    }
    int k = t[s].size();
    for(int i = 0;i < k;i++) {
        if(sign[t[s][i]] == 0) {
            dfs1(t[s][i],np,fi);
        }
    }
    sign[s] = 0;
    ape[col[s]]--;
    return 0;
}

int dfs2(int s,ll pr,int siz,int st,ll sn) {
    sign[s] = 1;
    ape[col[s]]++;
    sum[s] += pr+sn;
    if(ape[col[s]] == 1) {
        if(wh == 1) {
            
        }
        sum[s] += (st-siz)-cnt[col[s]];
        sn += (st-siz)-cnt[col[s]];
    }
    int k = t[s].size();
    for(int i = 0;i < k;i++) {
        if(sign[t[s][i]] == 0) {
            dfs2(t[s][i],pr,siz,st,sn);
        }
    }
    sign[s] = 0;
    ape[col[s]]--;
    return 0;
}

int divide(int s) {
    to = 0;
    wh = 0;
    ll chk;
    mi = 500000000;
    wt(s,size[s]);
    int w = wh;
    up_size(w);
    
    chk = sum[w];//基本跨区间贡献
    dfscore(w,w,0);
    chk = sum[w]-chk;

    
    int k = t[w].size();
    ape[col[w]] = 1;
    sign[w] = 1;
    for(int i = 0;i < k;i++) {
        if(sign[t[w][i]] == 0) {
            dfs1(t[w][i],0,w);
            cnt[col[w]] -= size[t[w][i]];
            chk = chk - to - size[t[w][i]];
            //printf("%d\n",to);
            dfs2(t[w][i],chk,size[t[w][i]],size[w],0);
            chk = chk + to + size[t[w][i]];
            cnt[col[w]] += size[t[w][i]];
            dfs1(t[w][i],1,w);
        }
    }
    
    
    
    ape[col[w]] = 0;
    for(int i = 1;i <= sz;i++) {
        cnt[lt[i]] = 0; ext[lt[i]] = 0;
    }
    sz = 0;
    
    if(w == 2) {
        
    }
    
    for(int i = 0;i < k;i++) {
        if(sign[t[w][i]] == 0) {
            divide(t[w][i]);
        }
    }
    
    return 0;
}



int main () {
    sign[0] = 1;
    scanf("%d",&n);
    for(int i = 1;i <= n;i++) {
        scanf("%d",&col[i]);
    }
    for(int i = 1;i <= n-1;i++) {
        int a,b;
        scanf("%d%d",&a,&b);
        t[a].push_back(b);
        t[b].push_back(a);
    }
    sign[0] = 1;
    up_size(1);
    for(int i = 1;i <= n;i++) {
      //  printf("%d ",size[i]);
    }
    divide(1);
    for(int i = 1;i <= n;i++) {
        printf("%lld\n",sum[i]);
    }
    return 0;
}