题目的意思是：给定一个点带颜色的树，两点之间的距离定义为路径上不同颜色的个数。求所有点对间的距离和。
做法有点分治，还有传说中的虚树DP，树上差分。
点分治法：
考虑每个点的贡献，可以发现一个点的子树大小就是这个点的贡献。那么，对于同一个根的另一个子树的一个点x，去掉x到根结点对应颜色的贡献，再加上x到根结点上的颜色的种类数目，就是这个x点的答案。我们具体做的时候，是先不考虑根结点的，根结点对x点的贡献单独算。
#include <algorithm> #include <iterator> #include <iostream> #include <cstring> #include <cstdlib> #include <iomanip> #include <bitset> #include <cctype> #include <cstdio> #include <string> #include <vector> #include <stack> #include <cmath> #include <queue> #include <list> #include <map> #include <set> #include <cassert> /* ⊂_ヽ ＼＼ Λ＿Λ 来了老弟 ＼(‘ㅅ‘) > ⌒ヽ / へ＼ / / ＼＼ ﾚ ノ ヽ_つ / / / /| ( (ヽ | |、＼ | 丿 ＼ ⌒) | | ) / ‘ノ ) Lﾉ */ using namespace std; #define lson (l , mid , rt << 1) #define rson (mid + 1 , r , rt << 1 | 1) #define debug(x) cerr << #x << " = " << x << "\n"; #define pb push_back #define pq priority_queue typedef long long ll; typedef unsigned long long ull; //typedef __int128 bll; typedef pair<ll ,ll > pll; typedef pair<int ,int > pii; typedef pair<int,pii> p3; //priority_queue<int> q;//这是一个大根堆q //priority_queue<int,vector<int>,greater<int> >q;//这是一个小根堆q #define fi first #define se second //#define endl ‘\n‘ #define boost ios::sync_with_stdio(false);cin.tie(0) #define rep(a, b, c) for(int a = (b); a <= (c); ++ a) #define max3(a,b,c) max(max(a,b), c); #define min3(a,b,c) min(min(a,b), c); const ll oo = 1ll<<17; const ll mos = 0x7FFFFFFF; //2147483647 const ll nmos = 0x80000000; //-2147483648 const int inf = 0x3f3f3f3f; const ll inff = 0x3f3f3f3f3f3f3f3f; //18 const int mod = 1e9+7; const double esp = 1e-8; const double PI=acos(-1.0); const double PHI=0.61803399; //黄金分割点 const double tPHI=0.38196601; template<typename T> inline T read(T&x){ x=0;int f=0;char ch=getchar(); while (ch<‘0‘||ch>‘9‘) f|=(ch==‘-‘),ch=getchar(); while (ch>=‘0‘&&ch<=‘9‘) x=x*10+ch-‘0‘,ch=getchar(); return x=f?-x:x; } inline void cmax(int &x,int y){if(x<y)x=y;} inline void cmax(ll &x,ll y){if(x<y)x=y;} inline void cmin(int &x,int y){if(x>y)x=y;} inline void cmin(ll &x,ll y){if(x>y)x=y;} /*-----------------------showtime----------------------*/ const int maxn = 2e5+9; int col[maxn]; vector<int>mp[maxn]; ll ans = 0, sumcol = 0; int sz[maxn],wt[maxn], root, curn; int vis[maxn]; void findRoot(int u, int fa) { sz[u] = 1;wt[u] = 0; for(int i=0; i<mp[u].size(); i++) { int v = mp[u][i]; if(v == fa || vis[v]) continue; findRoot(v, u); sz[u] += sz[v]; wt[u] = max(sz[v], wt[u]); } wt[u] = max(wt[u], curn - sz[u]); if(wt[u] <= wt[root]) root = u; } // map<int, int> pp; ll pp[maxn]; int youmeiyou[maxn]; int ss; void gao(int u, int fa, vector<pii>& vv, int cnt, ll sumfa, ll sum) { ll res = 0; if(youmeiyou[col[u]] == 0) vv.pb(pii(col[u], sz[u])), cnt++, res += pp[col[u]]; youmeiyou[col[u]]++; ans += sumcol - sumfa - res + 1ll * cnt * sum; //sum-color[根的颜色]+size[root] if(youmeiyou[col[ss]] == 0) ans += sum - pp[col[ss]]; for(int i=0; i<mp[u].size(); i++) { int v = mp[u][i]; if(fa == v || vis[v]) continue; gao(v, u, vv, cnt, sumfa + res, sum); } youmeiyou[col[u]] --; } void solve(int u) { vis[u] = 1; findRoot(u, -1); ll sum = 1; sumcol = 0; queue<int>needclear; needclear.push(col[u]); for(int i=0; i<mp[u].size(); i++) { int v = mp[u][i]; if(vis[v]) continue; vector<pii>vv; ss = u; gao(v, -1, vv, 0, 0, sum); for(int j=0; j<vv.size(); j++){ int c = vv[j].fi; if(pp[c])pp[c] += vv[j].se; else { pp[c] = vv[j].se; needclear.push(c); } sumcol += vv[j].se; } sum += sz[v]; } while(!needclear.empty()) { pp[needclear.front()] = 0; needclear.pop(); } for(int i=0; i<mp[u].size(); i++) { int v = mp[u][i]; if(!vis[v]) { root = 0; wt[0] = inf; curn = sz[v]; findRoot(v, -1); solve(root); } } } int main(){ int n, cas = 0; while(~scanf("%d", &n)) { memset(vis, 0, sizeof(vis)); for(int i=1; i<=n; i++) scanf("%d", &col[i]); for(int i=1; i<=n; i++) mp[i].clear(); for(int i=1; i<n; i++) { int u,v; scanf("%d%d", &u, &v); mp[u].pb(v); mp[v].pb(u); } ans = 0; root = 0; wt[0] = inf; curn = n; findRoot(1, -1); solve(root); printf("Case #%d: %lld\n", ++cas, ans); } return 0; } /* 6 1 2 3 1 2 3 1 2 1 3 3 4 3 5 4 6 */
原文：https://www.cnblogs.com/ckxkexing/p/11129428.html
