题解 P2664 【树上游戏】
「QQ红包」
2018-04-09 14:22:21
[my blog:redbag的小屋](http://redbag.pw/)
$$ \sum _{i=1}^n s(i,j) $$
考虑离线，一种一种颜色加进来。
i到没有出现这种颜色的地方，构成一个联通块。(用并查集来维护) $ ans+=(n-size) $
复杂度最坏 $ O(n^2) $ 。
期望得分 40分。
然后想想怎么把40分变成一百分。
每次点分的时候只要算子树对子树的贡献，还有子树对根的贡献，内部贡献继续分的时候算。
每个子树的颜色x对其他子树的贡献=该子树所有链上，x第一次出现的位置p，p的子树的大小。然后遇到受到内部影响的情况，贡献就要修改。
多种颜色我们不可以分开处理，就一起处理（看着鬼畜）
1.找重心，求size,清上一次的标记
2.第一次dfs：记录整个子树，每个颜色在该条链上第一次出现时候，其子树的大小的和。
所以这到底是什么东西啊QAQ。
看图感性理解写（逃），
然后还要记录一下每一种颜色，在整个分治块的答案，
然后和根颜色不同的才算。记一个总和s2，每个树的贡献就是sum[i]。因为根的颜色的贡献就是整个分治块的大小（想一想，为什么）。
3.第二次dfs：设整个子树对其他子树和根的贡献为sum[i]，设s1=s2-sum[i]-size[i]（算自己的时候不要算内部的贡献），
如果出现了颜色z（并且z ≠根的颜色），颜色z在整个分治块的总贡献为cnt[z], 目前在点d，所在子树为x，
cnt1[z]表示颜色z在目前子树x的答案（第三次dfs），s1=s1-cnt[z]+cnt1[c[x]]+（sz-size[x]）
（就是把原本这种颜色的贡献改为整个分治块的大小减去自己整个子树的大小）
4.然后还需要一次dfs来清理标记QAQ。
语文不好不能表达得很清楚，代码也不是很清晰，有问题可以评论/私信。
```cpp
#include<bits/stdc++.h>
#define ld long double
#define ll long long
using namespace std;
void qmax(int &x,int y) {if (x<y) x=y;}
void qmin(int &x,int y) {if (x>y) x=y;}
inline int read()
{
char s;
int k=0,base=1;
while((s=getchar())!='-'&&s!=EOF&&!(isdigit(s)));
if(s==EOF)exit(0);
if(s=='-')base=-1,s=getchar();
while(isdigit(s)) {k=k*10+(s^'0');s=getchar();}
return k*base;
}
inline void write(int x)
{
static char cnt,num[15];cnt=0;
if (!x)
{
printf("0");
return;
}
for (;x;x/=10) num[++cnt]=x%10;
for (;cnt;putchar(num[cnt--]+48));
}
const int maxn=1e5+100;
int n,flag;
int c[maxn];//颜色
int to[maxn<<1],ne[maxn<<1],po[maxn],id,X,Y;
int size[maxn],hson[maxn],root,del[maxn],sz;//del删除标记
ll ans[maxn];//答案
ll sum[maxn],s1;//记录每个子树的答案
int vis[maxn];//标记某种颜色是否在该链中出现过
ll cnt[maxn]; //存整个分治块的答案
int cnt1[maxn];//存当前扫的子树的答案
void add(int x,int y)
{
id++;to[id]=y;ne[id]=po[x];po[x]=id;
}
void get_root(int x,int fa)
{
cnt[c[x]]=0;//清理标记
size[x]=1;hson[x]=0;
for (int i=po[x];i;i=ne[i])
{
if (!del[to[i]]&&fa!=to[i])
{
get_root(to[i],x);
if (size[to[i]]>hson[x]) hson[x]=size[to[i]];
size[x]+=size[to[i]];
}
}
qmax(hson[x],sz-size[x]);
if (root==-1||hson[root]>hson[x]) root=x;
}//找重心
void qsize(int x,int fa)
{
size[x]=1;
for (int i=po[x];i;i=ne[i])
{
if (!del[to[i]]&&fa!=to[i])
{
qsize(to[i],x);
size[x]+=size[to[i]];
}
}
}//求每个子树的大小
void dfs1(int x,int fa,int s)//目前在点x,父亲为fa，所属于子树s
{
vis[c[x]]++; //标记出现次数
if (vis[c[x]]==1&&c[x]!=c[root])
{
cnt[c[x]]+=size[x];
sum[s]+=size[x];//该分治块的贡献
}
for (int i=po[x];i;i=ne[i])
{
if (to[i]==fa||del[to[i]]) continue;
dfs1(to[i],x,s);
}
vis[c[x]]--;
}
void dfs3(int x,int fa,int s)
{
vis[c[x]]++;
if (vis[c[x]]==1&&c[x]!=c[root]) cnt1[c[x]]+=size[x];// 该子树的贡献
for (int i=po[x];i;i=ne[i])
{
if (to[i]==fa||del[to[i]]) continue;
dfs3(to[i],x,s);
}
vis[c[x]]--;//清标记
}
void dfs4(int x,int fa,int s)
{
cnt1[c[x]]=0;
for (int i=po[x];i;i=ne[i])
{
if (to[i]==fa||del[to[i]]) continue;
dfs4(to[i],x,s);
}
}//清标记而已233
void dfs2(int x,int fa,int s)
{
vis[c[x]]++;
ll s3;
if (c[x]!=c[root]&&vis[c[x]]==1)//修改
{
s3=s1;
s1=s1-cnt[c[x]]+cnt1[c[x]]+sz-size[s];
}
ans[x]+=s1;
for (int i=po[x];i;i=ne[i])
{
if (to[i]==fa||del[to[i]]) continue;
dfs2(to[i],x,s);
}
if (c[x]!=c[root]&&vis[c[x]]==1) s1=s3;//回溯掉qaq
vis[c[x]]--;
}
void solve(int x)
{
qsize(x,0);//求出每个树的size
if (size[x]==1)
{
ans[x]++;//自己到自己也要算
del[x]=1;
return;
}
ll s2=0;
for (int i=po[x];i;i=ne[i])
{
if (del[to[i]]) continue;
sum[to[i]]=0;
dfs1(to[i],x,to[i]);
s2+=sum[to[i]];
}
s2+=size[root];
ans[root]+=s2;
for (int i=po[x];i;i=ne[i])
{
if (del[to[i]]) continue;
s1=s2-sum[to[i]]-size[to[i]];
dfs3(to[i],x,to[i]);
dfs2(to[i],x,to[i]);
dfs4(to[i],x,to[i]);
} //统答案
del[x]=1;
for (int i=po[x];i;i=ne[i])
{
if (del[to[i]]) continue;
sz=size[to[i]];
root=-1;
get_root(to[i],x);
solve(root);
}//处理下一层
}
int main()
{
n=read();
for (int i=1;i<=n;i++) c[i]=read();
for (int i=1;i<n;i++)
{
X=read();Y=read();
add(X,Y),add(Y,X);
}
sz=n;
root=-1;
get_root(1,0);
solve(root);
for (int i=1;i<=n;i++)
{
printf("%lld\n",ans[i]);
}
return 0;
}
```