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子树的路径，所以一开始要先减掉。

代码：

  1 #include<cstdio>
  2 #include<algorithm>
  3 #include<cstring>
  4 #include<cmath>
  5 #include<iostream>
  6 #include<cctype>
  7 #include<set>
  8 #include<vector>
  9 #include<queue>
 10 #include<map>
 11 #define fi(s) freopen(s,"r",stdin);
 12 #define fo(s) freopen(s,"w",stdout);
 13 using namespace std;
 14 typedef long long LL;
 15 
 16 inline int read() {
 17     int x=0,f=1;char ch=getchar();for(;!isdigit(ch);ch=getchar())if(ch=='-')f=-1;
 18     for(;isdigit(ch);ch=getchar())x=x*10+ch-'0';return x*f;
 19 }
 20 
 21 
 22 const int INF = 1e9;
 23 const int N = 100005;
 24 
 25 int head[N], nxt[N << 1], to[N << 1], En;
 26 int col[N], siz[N], sk[N];
 27 LL ans[N], cnt[N], sum[N], Sum, Num;
 28 bool vis[N];
 29 int Size, Mn, Root, top;
 30 // cnt[i] 颜色i出现的次数，sum[i]以根为起点，包含颜色i的路径条数，Sum为sum[i]的和。 
 31 
 32 void add_edge(int u,int v) {
 33     ++En; to[En] = v; nxt[En] = head[u]; head[u] = En;
 34     ++En; to[En] = u; nxt[En] = head[v]; head[v] = En;
 35 }
 36 
 37 void getroot(int u,int fa) {
 38     siz[u] = 1;
 39     int cnt = 0;
 40     for (int i=head[u]; i; i=nxt[i]) {
 41         int v = to[i];
 42         if (v == fa || vis[v]) continue; // vis[v]!!!
 43         getroot(v, u);
 44         siz[u] += siz[v];
 45         cnt = max(cnt, siz[v]);
 46     }
 47     cnt = max(cnt, Size - siz[u]);
 48     if (cnt < Mn) { Mn = cnt, Root = u; }
 49 }
 50 
 51 void dfs1(int u,int fa) { // 计算siz，sum，以根为起点的答案。 
 52     siz[u] = 1; cnt[col[u]] ++;
 53     for (int i=head[u]; i; i=nxt[i]) 
 54         if (!vis[to[i]] && to[i] != fa) 
 55             dfs1(to[i], u), siz[u] += siz[to[i]];
 56     if (cnt[col[u]] == 1) 
 57         Sum += siz[u], sum[col[u]] += siz[u], sk[++top] = col[u];
 58     cnt[col[u]] --;
 59 }
 60 
 61 void dfs2(int u,int fa) { // 计算子树内每个点 过根的所有路径的答案。 
 62     cnt[col[u]] ++;
 63     if (cnt[col[u]] == 1) 
 64         Num ++, Sum -= sum[col[u]]; // 只考虑过根的路径，Num记录这个点到根的路径第一次出现的颜色的个数 
 65     ans[u] += Num * Size + Sum; 
 66     // 这些Num个颜色因为到根的路径上已经有这个颜色了，所以和其他的点任意组合的路径都满足有这个颜色；Sum为除了Num个颜色以外的颜色的贡献 
 67     for (int i=head[u]; i; i=nxt[i]) 
 68         if (!vis[to[i]] && to[i] != fa) dfs2(to[i], u);
 69     if (cnt[col[u]] == 1) 
 70         Num --, Sum += sum[col[u]];
 71     cnt[col[u]] --;
 72 }
 73 
 74 void Modify(int u,int fa,int p) {
 75     cnt[col[u]] ++;
 76     for (int i=head[u]; i; i=nxt[i]) 
 77         if (!vis[to[i]] && to[i] != fa) Modify(to[i], u, p);
 78     if (cnt[col[u]] == 1) 
 79         Sum += siz[u] * p, sum[col[u]] += siz[u] * p;
 80     cnt[col[u]] --;
 81 }
 82 void change(int u,int fa,int p) {
 83     Sum += siz[u] * p, sum[col[fa]] += siz[u] * p; // sum[col[fa]]=siz[fa]，现在应该减去这棵子树 
 84     cnt[col[fa]] ++; Modify(u, fa, p); cnt[col[fa]] --;
 85 }
 86 
 87 void Calc(int u) {
 88     top = 0; dfs1(u, 0); ans[u] += Sum; // 计算子树内每个点的siz，sum，以根为起点的答案。 
 89     for (int i=head[u]; i; i=nxt[i]) {
 90         int v = to[i];
 91         if (vis[v]) continue;
 92         change(v, u, -1); // 把这棵子树的贡献减去 
 93         Size = siz[u] - siz[v]; dfs2(v, u); // Size = siz[v] !!!，计算子树内的每个点答案。 
 94         change(v, u, 1); // 加回来 
 95     }
 96     Num = Sum = 0;
 97     for (int i=1; i<=top; ++i) cnt[sk[i]] = sum[sk[i]] = 0;
 98 }
 99 void solve(int u) {
100     Calc(u); vis[u] = true;
101     for (int i=head[u]; i; i=nxt[i]) {
102         int v = to[i];
103         if (vis[v]) continue;
104         Size = siz[v], Mn = INF;
105         getroot(v, 0);
106         solve(Root);
107     }
108 }
109 
110 int main() { 
111     int n = read();
112     for (int i=1; i<=n; ++i) col[i] = read();
113     for (int i=1; i<n; ++i) {
114         int x = read(), y = read();
115         add_edge(x, y);
116     }
117     Size = n, Mn = 1e9;
118     getroot(1, 0);
119     solve(Root);
120     for (int i=1; i<=n; ++i) printf("%lld
",ans[i]);
121     return 0;
122 }