Problem

https://www.luogu.org/problemnew/show/P2664

Solution

论打标记的艺术？？

https://www.luogu.org/blog/user20176/solution-p2664

这篇题解写的很清楚了。

主要是实现的问题。

那么我是先把所有的子树处理好后把重儿子存下来，以便下一步访问。

然后在递归每个子树对其更新 sum 值。在递归每个子树之前先递归求其子树对 cnt 的贡献 cnt2。然后相减就行了。

那你说不是要做很多遍 memset 使复杂度退化？

这就是打标记的艺术了。

顺便，少数点会爆 int。

Code

// Code by ajcxsu
// Problem: game on the tree

#include<bits/stdc++.h>
#define INF (0x3f3f3f3f)
using namespace std;
typedef long long ll;

inline void gn(int &x) {
    char ch=getchar();
    x=0;
    while(ch<'0' || ch>'9') ch=getchar();
    while(ch>='0' && ch<='9') x=x*10+ch-'0', ch=getchar();
}

const int N=1e5+10, M=2e5+10, P=20;
int h[N], to[M], nexp[M], p=1;
inline void ins(int a, int b) { nexp[p]=h[a], h[a]=p, to[p]=b, p++; }

int n;
int c[N]; // color

int S, hvy, minn, siz[N];
ll sum[N]; // 记录答案
ll cnt[N]; // 记录颜色的贡献
ll cnt2[N];
ll tot; // 颜色贡献的总和
ll tot2;
int app[N]; // 标记颜色是否出现过（用于回溯）
int app2[N]; // 用于分治中心
int app3[N], tag;
vector<int> son[N]; // 重儿子记录
int size; // 记录分治中心子树


bool vis[N];

void ghvy(int x, int fa) {
    siz[x]=1;
    int ma=0;
    for(int u=h[x];u;u=nexp[u])
        if(to[u]!=fa && !vis[to[u]]) {
            ghvy(to[u], x);
            siz[x]+=siz[to[u]];
            ma=max(ma, siz[to[u]]);
        }
    ma=max(ma, S-siz[x]);
    if(ma<minn) minn=ma, hvy=x;
}

void gvan(int x, int fa, int p, int idx) {
    bool rev=false;
    if(app[c[x]]!=idx) {
        if(app2[c[x]]!=p) {
            app2[c[x]]=p, cnt[c[x]]=0;
        }
        cnt[c[x]]+=siz[x], tot+=siz[x], app[c[x]]=idx;
        rev=1;
    }
    for(int u=h[x];u;u=nexp[u])
        if(to[u]!=fa && !vis[to[u]]) {
            gvan(to[u], x, p, idx);
        }
    if(rev) app[c[x]]=0; // 回溯
}

void gvans(int x, int fa, int p, int idx) {
    bool rev=false;
    if(app[c[x]]!=idx) {
        if(app3[c[x]]!=tag) {
            app3[c[x]]=tag, cnt2[c[x]]=0;
        }
        cnt2[c[x]]+=siz[x], tot2+=siz[x], app[c[x]]=idx;
        rev=1;
    }
    for(int u=h[x];u;u=nexp[u])
        if(to[u]!=fa && !vis[to[u]]) {
            gvans(to[u], x, p, idx);
        }
    if(rev) app[c[x]]=0; // 回溯
}

void gsum(int x, int fa, int p, int idx, ll val) {
    bool rev=false;
    if(app[c[x]]!=idx) {
        if(app3[c[x]]!=tag) cnt2[c[x]]=0;
        val+=size-siz[idx]-(cnt[c[x]]-cnt2[c[x]]);
        app[c[x]]=idx;
        rev=1;
    }
    sum[x]+=val;
    for(int u=h[x];u;u=nexp[u])
        if(to[u]!=fa && !vis[to[u]]) {
            gsum(to[u], x, p, idx, val);
        }
    if(rev) app[c[x]]=0;
}

void gans(int x) {
    vis[x]=1;
    size=1, tot=0;
    for(int u=h[x];u;u=nexp[u])
        if(!vis[to[u]]) {
            ghvy(to[u],0);

            app[c[x]]=to[u];
            gvan(to[u], 0, x, to[u]);
            app[c[x]]=0;
            size+=siz[to[u]];

            S=siz[to[u]], minn=INF, ghvy(to[u],0);
            son[x].push_back(hvy);
        }
    sum[x]+=size+tot;
    for(int u=h[x];u;u=nexp[u])
        if(!vis[to[u]]) {
            ++tag;
            app[c[x]]=to[u];
            tot2=0;
            gvans(to[u], 0, x, to[u]);
            gsum(to[u], 0, x, to[u], tot-tot2+size-siz[to[u]]);
            app[c[x]]=0;
        }
    for(int i=0, len=son[x].size();i<len;i++)
        gans(son[x][i]);
}


int main() {
    gn(n);
    for(int i=1;i<=n;i++) gn(c[i]);
    for(int u,v,i=0;i<n-1;i++) gn(u), gn(v), ins(u,v), ins(v,u);
    gans(1);
    for(int i=1;i<=n;i++) printf("%lld\n", sum[i]);
    return 0;
}