2019年ccpc女生赛重现赛题解E
题目：
Checkout
Time Limit: 2000/1000 MS (Java/Others) Memory Limit: 524288/524288 K (Java/Others)
Total Submission(s): 0 Accepted Submission(s): 0
Problem Description
Alice 是一个身怀改变世界的抱负的著名企业家，手中掌控很多著名的公司，为了更好的管理， Alice 建立了一套很完善的架构体系，已知 Alice 的企业的架构体系是一棵树，每个节点代表一个人。对于每个节点，它的父节点就是这个人的直接 leader，每个节点都有一个权值，代表这个人的爱好，每对属于同一直接 leader 的节点如果有相同的爱好那么他们就会给公司产生1的和谐度，如果一个人和他/她的直接 leader 的爱好相同也会给公司产生1的和谐度，但是再完美的公司都会产生离职的情况，如果一个人离职，并且这个人有直接下属，就从这个人的直接下属中选一个成为这个部门新的 leader，使得新的组织架构的和谐度最高，新晋的人汇报给离职的人的 leader（如果存在的话），同时新 leader 的原来下属的直接 leader 不变， Alice 想知道如果某个人离职整个公司的和谐度是多少。
Input
第一行两个正整数 n, m 分别代表公司的总人数和爱好的种数。
第二行有 n 个整数，第 i 个数 ai，代表第 i 个人的爱好。
后面 n − 1 行每行两个数u，v代表 u，v 存在上下属关系（上下级关系不确定）。
假设树根为 1，即 1 号员工是公司的ceo，他/她不需要汇报给任何人。
1 ≤ m ≤ n ≤ 100, 000
1 ≤ ai ≤ m
Output
输出一行n个数，第i个数代表如果编号为 i 的人离职，公司的和谐度会是多少，以空格分隔，行末无空格。
Sample Input
4 3
2 3 1 3
1 2
2 3
2 4
Sample Output
1 0 1 0
思路：树形dp。
AC代码：
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=1e5+5;
int vis[maxn],Fa[maxn];
ll sum[maxn],w[maxn],a[maxn],color[maxn],fac[maxn];
vector<int>g[maxn];
void dfs(int u,int fa){
sum[u]=0;
Fa[u]=fa;
w[u]=0;
for(int i=0;i<g[u].size();i++){
int v=g[u][i];
if(v==fa) continue;
dfs(v,u);
sum[u]+=sum[v];
}
for(int i=0;i<g[u].size();i++){//将 该节点的儿子节点颜色加入桶中
int v=g[u][i];
if(v==fa) continue;
color[a[v]]++;
}
w[u]=color[a[u]];//该点与下属的贡献即为桶中 该点颜色的数量
//统计儿子节点之间产生的贡献
vector<int>q;//存用了哪几种颜色
for(int i=0;i<g[u].size();i++){
int v=g[u][i];
if(v==fa) continue;
if(!vis[a[v]])//如果这个点颜色没用过
{
q.push_back(a[v]);//把颜色丢进去
sum[u]+=(color[a[v]]-1)*(color[a[v]])/2;
vis[a[v]]=1;
}
}
//
for(int i=0;i<q.size();i++) vis[q[i]]=0;//清空vis数组
for(int i=0;i<g[u].size();i++){//将 该节点的儿子节点颜色删除 清空桶
int v=g[u][i];
if(v==fa) continue;
color[a[v]]--;
}
sum[u]+=w[u];
}
ll ans[maxn];
int main()
{
ios::sync_with_stdio(false);
int n,m,u,v;
cin>>n>>m;
for(int i=1;i<=n;i++) cin>>a[i];
for(int i=1;i<n;i++){
cin>>u>>v;
g[u].push_back(v);
g[v].push_back(u);
}
g[0].push_back(1);
dfs(1,0);
for(int u=1;u<=n;u++){//枚举删除那个点
ans[u]=sum[1]-w[u];//贡献必然先减去离职点的贡献
if(a[u]==a[Fa[u]])ans[u]--;//如果离职点和leader颜色相同 贡献减一
//cout<<"U: "<<u<<". 1: "<<ans[u]<<endl;
// 离职点同层点的情况
for(int i=0;i<g[Fa[u]].size();i++){
int v=g[Fa[u]][i];
if(v==Fa[v]) continue;
fac[a[v]]++;//该层点颜色放入桶中
}
ans[u]-=fac[a[u]]-1;
fac[a[u]]--;//离职点颜色删除
//离职点的儿子节点颜色插入桶中
for(int i=0;i<g[u].size();i++){
int v=g[u][i];
if(v==Fa[u]) continue;
color[a[v]]++;
}
//cout<<"U: "<<u<<". 2: "<<ans[u]<<endl;
ll temp=-1e18;//记录改变值最小
for(int i=0;i<g[u].size();i++){//枚举离职点哪位儿子继位
ll res=0;//记录以该点时的贡献
int v=g[u][i];
if(v==Fa[u]) continue;
if(a[v]==a[Fa[u]]) res++;// 升职后和新的leader颜色相同 贡献加1
res-=color[a[v]]-1;// 该点升职原本为兄弟时的贡献删去
//计算儿子变兄弟 产生的额外贡献
for(int j=0;j<g[v].size();j++){
int vv=g[v][j];
if(vv==Fa[v]) continue;
res+=fac[a[vv]];
}
temp=max(temp,res);
//该儿子贡献计算完
}
//离职点的儿子节点颜色删除
for(int i=0;i<g[u].size();i++){
int v=g[u][i];
if(v==Fa[u]) continue;
color[a[v]]--;
}
//删离职同层的颜色
for(int i=0;i<g[Fa[u]].size();i++){//删边
int v=g[Fa[u]][i];
if(v==Fa[v]) continue;
fac[a[v]]--;//删去
}
fac[a[u]]++;//离职点删除多于一次补回
//cout<<"U: "<<u<<". 3: "<<ans[u]<<' '<<temp<<endl;
if(temp==-1e18) continue;
ans[u]+=temp;
//cout<<"U: "<<u<<".final ans: "<<ans[u]<<endl;
}
// for(int i=1;i<=n;i++) cout<<w[i]<<' ';cout<<endl;
// for(int i=1;i<=n;i++) cout<<sum[i]<<' ';cout<<endl;
// for(int i=1;i<=n;i++) cout<<color[i]<<' '<<fac[i]<<' ';cout<<endl;
for(int i=1;i<=n-1;i++){
cout<<ans[i]<<' ';
}
cout<<ans[n]<<endl;
return 0;
}