P2664 树上游戏

P2664 树上游戏

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

分析：

　　点分治。

　　首先关于答案的统计转化成计算每个颜色的贡献。

　　1、计算从根出发的路径的答案：如果某一个颜色是从根到这个点的链上的第一次出现的，那么这个颜色会对根产生siz[x]个贡献。（根连向它子树的任意一个点的路径都包含这个颜色）。

　　2、计算子树内每个点过根的路径答案：记录一个数组sum[i]，表示从根出发包含颜色i的路径的条数（在1中，找到一个第一次出现的颜色，加上这个点的siz即可）。然后假设当前点是x，根为z，x所在的子树为y。x->z的路径上，出现的颜色为Num，那么这Num个颜色由于已经在到根的路径上有了，那么随便选择其它子树内的点构成的路径都包含了这个颜色，贡献为(siz[z]-siz[y])*Num；未出现的颜色的贡献：在y子树外计算多少个点与x构成的路径，包含这个颜色。那么可以sum数组的和得到，但是其中包含的y子树的路径，所以一开始要先减掉。

代码：

 #include<cstdio>
 #include<algorithm>
 #include<cstring>
 #include<cmath>
 #include<iostream>
 #include<cctype>
 #include<set>
 #include<vector>
 #include<queue>
 #include<map>
 #define fi(s) freopen(s,"r",stdin);
 #define fo(s) freopen(s,"w",stdout);
 using namespace std;
 typedef long long LL;

 inline int read() {
     int x=,f=;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-;
     for(;isdigit(ch);ch=getchar())x=x*+ch-'';return x*f;
 }

 const int INF = 1e9;
 const int N = ;

 int head[N], nxt[N << ], to[N << ], En;
 int col[N], siz[N], sk[N];
 LL ans[N], cnt[N], sum[N], Sum, Num;
 bool vis[N];
 int Size, Mn, Root, top;
 // cnt[i] 颜色i出现的次数，sum[i]以根为起点，包含颜色i的路径条数，Sum为sum[i]的和。 

 void add_edge(int u,int v) {
     ++En; to[En] = v; nxt[En] = head[u]; head[u] = En;
     ++En; to[En] = u; nxt[En] = head[v]; head[v] = En;
 }

 void getroot(int u,int fa) {
     siz[u] = ;
     int cnt = ;
     for (int i=head[u]; i; i=nxt[i]) {
         int v = to[i];
         if (v == fa || vis[v]) continue; // vis[v]!!!
         getroot(v, u);
         siz[u] += siz[v];
         cnt = max(cnt, siz[v]);
     }
     cnt = max(cnt, Size - siz[u]);
     if (cnt < Mn) { Mn = cnt, Root = u; }
 }

 void dfs1(int u,int fa) { // 计算siz，sum，以根为起点的答案。
     siz[u] = ; cnt[col[u]] ++;
     for (int i=head[u]; i; i=nxt[i])
         if (!vis[to[i]] && to[i] != fa)
             dfs1(to[i], u), siz[u] += siz[to[i]];
     if (cnt[col[u]] == )
         Sum += siz[u], sum[col[u]] += siz[u], sk[++top] = col[u];
     cnt[col[u]] --;
 }

 void dfs2(int u,int fa) { // 计算子树内每个点 过根的所有路径的答案。
     cnt[col[u]] ++;
     if (cnt[col[u]] == )
         Num ++, Sum -= sum[col[u]]; // 只考虑过根的路径，Num记录这个点到根的路径第一次出现的颜色的个数
     ans[u] += Num * Size + Sum;
     // 这些Num个颜色因为到根的路径上已经有这个颜色了，所以和其他的点任意组合的路径都满足有这个颜色；Sum为除了Num个颜色以外的颜色的贡献
     for (int i=head[u]; i; i=nxt[i])
         if (!vis[to[i]] && to[i] != fa) dfs2(to[i], u);
     if (cnt[col[u]] == )
         Num --, Sum += sum[col[u]];
     cnt[col[u]] --;
 }

 void Modify(int u,int fa,int p) {
     cnt[col[u]] ++;
     for (int i=head[u]; i; i=nxt[i])
         if (!vis[to[i]] && to[i] != fa) Modify(to[i], u, p);
     if (cnt[col[u]] == )
         Sum += siz[u] * p, sum[col[u]] += siz[u] * p;
     cnt[col[u]] --;
 }
 void change(int u,int fa,int p) {
     Sum += siz[u] * p, sum[col[fa]] += siz[u] * p; // sum[col[fa]]=siz[fa]，现在应该减去这棵子树
     cnt[col[fa]] ++; Modify(u, fa, p); cnt[col[fa]] --;
 }

 void Calc(int u) {
     top = ; dfs1(u, ); ans[u] += Sum; // 计算子树内每个点的siz，sum，以根为起点的答案。
     for (int i=head[u]; i; i=nxt[i]) {
         int v = to[i];
         if (vis[v]) continue;
         change(v, u, -); // 把这棵子树的贡献减去
         Size = siz[u] - siz[v]; dfs2(v, u); // Size = siz[v] !!!，计算子树内的每个点答案。
         change(v, u, ); // 加回来
     }
     Num = Sum = ;
     for (int i=; i<=top; ++i) cnt[sk[i]] = sum[sk[i]] = ;
 }
 void solve(int u) {
     Calc(u); vis[u] = true;
     for (int i=head[u]; i; i=nxt[i]) {
         int v = to[i];
         if (vis[v]) continue;
         Size = siz[v], Mn = INF;
         getroot(v, );
         solve(Root);
     }
 }

 int main() {
     int n = read();
     for (int i=; i<=n; ++i) col[i] = read();
     for (int i=; i<n; ++i) {
         int x = read(), y = read();
         add_edge(x, y);
     }
     Size = n, Mn = 1e9;
     getroot(, );
     solve(Root);
     for (int i=; i<=n; ++i) printf("%lld\n",ans[i]);
     return ;
 }