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子树的路径，所以一开始要先减掉。
代码：
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\n",ans[i]); 121 return 0; 122 }
