P2664 树上游戏 题解
原题链接：P2664。
题意：给定一棵树，每个点都有一个颜色 \(c_{i}\)。对于每一个点 \(i\)，求出 \(\sum_{j=1}^{n}s(i,j)\) 的值。其中 \(s(i,j)\) 表示点 \(i\) 到点 \(j\) 的颜色数量。
路径相关，考虑点分治。
假设当前的重心为 \(u\)，那么对 \(u\) 自己的贡献就是 \(u\) 到 \(u\) 的所有子树中每一个点的路径上的颜色数量之和，我们可以用 \(cnt_{i}\) 表示包含了颜色 \(i\) 的路径有多少条，则贡献为 \(tot=\sum cnt_{i}\)，直接 dfs 一遍即可，如果当前遍历到的点 \(v\) 的颜色在当前路径中是第一次出现，那么 \(cnt_{c_{v}}\) 应该加上 \(v\) 的子树大小（\(v\) 的子树中的每一个点到重心 \(u\) 的路径一定都包含 \(c_{v}\) 这个颜色），否则不加。这个颜色必须是第一次出现才加的原因是：一条路径上如果有多个相同颜色那么总共只会贡献 \(1\)。
然后考虑如何计算 \(u\) 的一颗子树中的一个节点的答案变化。我们可以将贡献拆分为两个部分。
-
其它子树中的任意一个点到重心 \(u\) 的贡献（路径前一半）
这一部分的贡献比较好求，可以用 \(u\) 的所有子树中的点到 \(u\) 的贡献 \(tot\) 减去当前 \(v\) 所在子树中的点到 \(u\) 的贡献。算当前子树的贡献的 dfs 方法和上述求 \(u\) 的贡献的方法相同。
-
重心 \(u\) 到当前子树的这个节点的贡献（路径后一半）
这一部分可以考虑对 \(v\) 所在的子树进行从上到下遍历，维护一个 \(tag\) 值表示当前的贡献，初始值为 \(0\)。如果当前遍历到的点 \(x\) 的颜色在路径中是第一次出现，那么 \(tag\) 需要加上 \(p-ct_{c_{x}}\)，其中 \(p\) 表示 \(u\) 到所有其它子树中的所有点的路径数，\(ct_{c_{x}}\) 表示 \(u\) 到所有其它子树中的点的路径中包含 \(c_{x}\) 这个颜色的路径数。当前的后一半路径一共可以匹配到 \(p\) 条前一半的路径，但是有 \(ct_{c_{x}}\) 条路径已经算了当前颜色 \(c_{x}\) 的贡献了，所以要减掉。
那么对于 \(u\) 的每一棵子树中的每一个点 \(v\)，答案应该加上 \((tot+tag\)），\(tot\) 是前一半的贡献，\(tag\) 是后一半的贡献。注意这里的 \(tot\) 应该减去 \(v\) 所对应的子树的大小 \(size\)，因为 \(c_{u}\) 的贡献多算了 \(size\) 遍，在算路径前一半的贡献的时候没有减掉。
dfs 是 \(O(n)\) 的，分治带一个 \(\log\)，总时间复杂度 \(O(n \log n)\)。
代码：
#include<bits/stdc++.h>
using namespace std;
#define int long
#define inf 0x3f
#define inf_db 127
#define ls id << 1
#define rs id << 1 | 1
#define re register
#define endl '\n'
typedef pair <int,int> pii;
const int MAXN = 1e5 + 10;
int n,sum[MAXN],head[MAXN],c[MAXN],x,y,Cnt,cnt[MAXN],color[MAXN],tp;
int size[MAXN],tot,top,s[MAXN],col[MAXN],ct[MAXN],cn[MAXN],path;
bool vis[MAXN],vi[MAXN];
struct Node{int u,v,nxt;}e[MAXN << 1];
inline void Add(int u,int v){e[++Cnt] = {u,v,head[u]};head[u] = Cnt;}
inline int Get_size(int u,int father)
{
if(vis[u] == true) return 0;
size[u] = 1;
for(int i = head[u]; ~ i;i = e[i].nxt)
if(e[i].v != father) size[u] += Get_size(e[i].v,u);
return size[u];
}
inline int Get_wc(int u,int father,int tot,int &wc)
{
if(vis[u] == true) return 0;
int sum = 1,mx = 0;
for(int i = head[u]; ~ i;i = e[i].nxt)
{
int now = e[i].v;
if(now == father) continue;
int tmp = Get_wc(now,u,tot,wc);
sum += tmp,mx = max(mx,tmp);
}
if(max(mx,tot - sum) <= tot / 2) wc = u;
return sum;
}
inline void Get_dist(int u,int father,int *cnt)
{
if(vis[u] == true) return;
if(vi[c[u]] == false) color[++top] = c[u],vi[c[u]] = true;
if(++s[c[u]] == 1) cnt[c[u]] += size[u];
for(int i = head[u]; ~ i;i = e[i].nxt)
{
int now = e[i].v;
if(now == father) continue;
Get_dist(now,u,cnt);
}
s[c[u]]--;
}
inline void Update(int u,int father,int tag)
{
if(vis[u] == true) return;
int val = tag;
if(++s[c[u]] == 1) val += path - cnt[c[u]];
sum[u] += tot + val;
for(int i = head[u]; ~ i;i = e[i].nxt)
{
int now = e[i].v;
if(now == father) continue;
Update(now,u,val);
}
s[c[u]]--;
}
inline void calc(int u)
{
if(vis[u] == true) return;
Get_wc(u,0,Get_size(u,0),u);
Get_size(u,0);
tot = top = 0,Get_dist(u,0,cnt);
vis[u] = true,tp = top;
for(int i = 1;i <= top;i++) vi[color[i]] = false;
for(int i = 1;i <= top;i++)
{
tot = tot + cnt[color[i]];
col[i] = color[i],cn[color[i]] = cnt[color[i]];
}
sum[u] += tot;int tmp = tot;
for(int i = head[u]; ~ i;i = e[i].nxt)
{
int now = e[i].v;
s[c[u]] = 1,top = 0;
Get_size(now,0);
Get_dist(now,u,ct),s[c[u]] = 0;
for(int j = 1;j <= top;j++) vi[color[j]] = false;
cnt[c[u]] -= size[now],tot -= size[now];
for(int j = 1;j <= top;j++) cnt[color[j]] -= ct[color[j]],tot -= ct[color[j]];
path = size[u] - size[now];
Update(now,u,0);
cnt[c[u]] += size[now],tot = tmp;
for(int j = 1;j <= top;j++) cnt[color[j]] = cn[color[j]],ct[color[j]] = 0;
}
for(int i = 1;i <= tp;i++) cnt[col[i]] = 0;
for(int i = head[u]; ~ i;i = e[i].nxt) calc(e[i].v);
}
signed main()
{
ios::sync_with_stdio(false);
cin.tie(0),cout.tie(0);
memset(head,-1,sizeof head);
cin >> n;
for(int i = 1;i <= n;i++) cin >> c[i];
for(int i = 1;i < n;i++) cin >> x >> y,Add(x,y),Add(y,x);
calc(1);
for(int i = 1;i <= n;i++) cout << sum[i] << endl;
return 0;
}
