可能是我有点菜，网上的题解都看不懂......但是写完之后看了好多题解代码，很多都有不影响正确性的思维上的小错误...
回顾下点分治的基理，就是选定重心，把重心的答案和一些点对经过重心路径产生的答案全部都处理出来，从而使得所有子树再也没有任何关系，从而独立而分治。
基本上O(nlogn)，树上点对问题，都是点分治跑不了。
我们考虑，我们选出重心后，我们显然可以用O(分块)的时间，处理出tot_ctb，表示所有点和根路径产生的答案。加到根的答案中。
我们再考虑，怎么处理一些点对经过重心路径产生的答案全部都处理出来，从而使得所有子树再也没有任何关系，从而独立而分治。我们考虑，我们对根节点求出的tot_ctb和ctb[]，如果能删去某一子树的影响，显然能够用来处理其他子树与该子树之间路径所产生的贡献。删去影响很简单，看代码即可。那怎么考虑产生的贡献呢，我们从子树的根开始搜到x，路上经过的颜色col，显然ctb[col]，即根节点的其他子树通过颜色col产生的贡献，不应在对x及其子树产生影响了，同时反观，每一种颜色，都会从x及其子树出发，到其余子树的任何一个点，产生贡献，sub表示其余子树中点的数量。
然后就ok啦，具体看代码。
1 #include <cstdio> 2 #include <algorithm> 3 using namespace std; 4 typedef long long ll; 5 const int MAXN = 101000,inf = 100000000; 6 int n,sum,maxn,cnt,root,num,sub; 7 int col[MAXN],ctb[MAXN],head[MAXN],siz[MAXN],tot[MAXN],to[MAXN << 1],nxt[MAXN << 1]; 8 ll res[MAXN]; 9 ll tot_ctb; 10 bool vis[MAXN]; 11 void add(int x,int y) 12 {//常规链前 13 nxt[++cnt] = head[x]; 14 to[cnt] = y; 15 head[x] = cnt; 16 } 17 void get_root(int x,int frm) 18 {//常规重心 19 siz[x] = 1; 20 int mx = 0; 21 for (int i = head[x];i;i = nxt[i]) 22 { 23 if (to[i] == frm || vis[to[i]] == true) 24 continue; 25 get_root(to[i],x); 26 siz[x] += siz[to[i]]; 27 mx = max(mx,siz[to[i]]); 28 } 29 mx = max(mx,sum - siz[x]); 30 if (mx < maxn) 31 { 32 maxn = mx; 33 root = x; 34 } 35 } 36 void dfs1(int x,int frm) 37 { 38 //维护出当前块的siz 39 //维护出tot_ctb，所有子树及根节点本身对根节点答案的贡献 40 //维护出ctb[col]，某种颜色col对根节点答案的贡献 41 siz[x] = 1; 42 tot[col[x]]++; 43 for (int i = head[x];i;i = nxt[i]) 44 { 45 if (to[i] == frm || vis[to[i]] == true) 46 continue; 47 dfs1(to[i],x); 48 siz[x] += siz[to[i]]; 49 } 50 if (tot[col[x]] == 1) 51 {//如果点x颜色是路径上第一次出现，显然x子树内，每个点到root都会因为x的颜色产生1贡献 52 tot_ctb += siz[x]; 53 ctb[col[x]] += siz[x]; 54 } 55 tot[col[x]]--; 56 } 57 void change(int x,int frm,int val) 58 {//清空某棵子树对ctb[],tot_ctb的影响 59 tot[col[x]]++; 60 for (int i = head[x];i;i = nxt[i]) 61 { 62 if (to[i] == frm || vis[to[i]] == true) 63 continue; 64 change(to[i],x,val); 65 } 66 if (tot[col[x]] == 1) 67 { 68 tot_ctb += siz[x] * val; 69 ctb[col[x]] += siz[x] * val; 70 } 71 tot[col[x]]--; 72 } 73 void dfs2(int x,int frm) 74 {//处理子树与当前块内其他点的所有影响，从而分治 75 tot[col[x]]++; 76 if (tot[col[x]] == 1) 77 {//如果某种颜色已经出现，那么显然这棵子树外的其他节点，不会对这里面的点答案产生贡献 78 tot_ctb -= ctb[col[x]]; 79 num++; 80 } 81 //num存放路径上出现颜色的个数，这些颜色均在子树外节点到当前点的路径上，产生了贡献 82 res[x] += tot_ctb + num * sub; 83 for (int i = head[x];i;i = nxt[i]) 84 { 85 if (to[i] == frm || vis[to[i]] == true) 86 continue; 87 dfs2(to[i],x); 88 } 89 if (tot[col[x]] == 1) 90 {//复原tot_ctb和num 91 tot_ctb += ctb[col[x]]; 92 num--; 93 } 94 tot[col[x]]--; 95 } 96 void clear(int x,int frm) 97 {//清空ctb[] 98 ctb[col[x]] = 0; 99 for (int i = head[x];i;i = nxt[i]) 100 { 101 if (to[i] == frm || vis[to[i]] == true) 102 continue; 103 clear(to[i],x); 104 } 105 } 106 107 void solve(int x) 108 {//注意这里的x是当前块的根 109 dfs1(x,0);//维护出tot_ctb,siz[],ctb[] 110 res[x] += tot_ctb;//我们直接把刚刚计算出的tot_ctb加到res[x]上 111 for (int i = head[x];i;i = nxt[i]) 112 { 113 if (vis[to[i]] == true) 114 continue; 115 //以下为了抹去子树to[i]，对tot_ctb和ctb[]的影响 116 tot[col[x]]++; 117 //在考虑根节点的情况下，定向搜索一颗子树做出的边界操作 118 //因为我们之前是dfs1(x,0),这次是 change(to[i],0),tot[col[x]]需要特殊处理 119 tot_ctb -= siz[to[i]]; 120 //x的颜色产生的贡献也要去掉 121 ctb[col[x]] -= siz[to[i]]; 122 change(to[i],x,-1); 123 tot[col[x]]--; 124 sub = siz[x] - siz[to[i]];//sub表示这棵子树会对多少个节点产生影响 125 dfs2(to[i],x);//将子树to[i]与当前块内其他点的影响处理完，以便后续分治 126 //以下为了恢复子树to[i],对tot_ctb和ctb[]的影响 127 tot[col[x]]++; 128 tot_ctb += siz[to[i]]; 129 ctb[col[x]] += siz[to[i]]; 130 change(to[i],x,1); 131 tot[col[x]]--; 132 } 133 tot_ctb = 0;num = 0;//复原全局变量，下次备用 134 clear(x,0);//清空相关的ctb[] 135 } 136 void pot_div(int x) 137 { 138 vis[x] = true; 139 solve(x); 140 for (int i = head[x];i;i = nxt[i]) 141 { 142 if (vis[to[i]] == true) 143 continue; 144 maxn = inf; 145 sum = siz[to[i]]; 146 get_root(to[i],x); 147 pot_div(root); 148 } 149 } 150 int main() 151 { 152 scanf("%d",&n); 153 for (int i = 1;i <= n;i++) 154 scanf("%d",&col[i]); 155 int tx,ty; 156 for (int i = 1;i <= n - 1;i++) 157 { 158 scanf("%d%d",&tx,&ty); 159 add(tx,ty); 160 add(ty,tx); 161 } 162 maxn = inf; 163 sum = n; 164 get_root(1,0); 165 pot_div(root); 166 for (int i = 1;i <= n;i++) 167 printf("%lld ",res[i]); 168 return 0; 169 }