dsu on tree(树上启发式合并)
/*树上启发式合并用于处理离线查询的子树问题第一次先求出重儿子第二次dfs的过程中,先计算轻儿子,贡献不保留,再计算重儿子,贡献保留最后再加上轻儿子的贡献,计算出当前节点的值复杂度为O(nlogn)*//*每个节点都有一个颜色编号计算一棵以1为根的树的以每个节点为根的子树中颜色最多的颜色编号和这题中子树的贡献就在于cnt数组,用来计数颜色出现了多少次*/#include <iostream&
·
/*
树上启发式合并
用于处理离线查询的子树问题
第一次先求出重儿子
第二次dfs的过程中,先计算轻儿子,贡献不保留,再计算重儿子,贡献保留
最后再加上轻儿子的贡献,计算出当前节点的值
复杂度为O(nlogn)
*/
/*
每个节点都有一个颜色编号
计算一棵以1为根的树的以每个节点为根的子树中颜色最多的颜色编号和
这题中子树的贡献就在于cnt数组,用来计数颜色出现了多少次
*/
#include <iostream>
#include <vector>
using namespace std;
const int maxn = 1e5 + 5;
typedef long long ll;
vector<int> g[maxn];
int c[maxn],cnt[maxn],son[maxn],siz[maxn];
ll ans[maxn];
void dfs1(int x,int f) //找重儿子
{
siz[x] = 1;
int maxsiz = 0;
for (int i = 0; i < g[x].size(); i++)
{
int t = g[x][i];
if( t == f ) continue;
dfs1(t,x);
siz[x] += siz[t];
if( maxsiz < siz[t] )
{
son[x] = t;
maxsiz = siz[t];
}
}
}
int flag = 0;
ll sum = 0,maxc = 0; //当前的答案和最大值
void count(int x,int f,int val) //将以x为根的子树的贡献加上val但不包括flag的贡献
{
cnt[c[x]] += val;
if( cnt[c[x]] > maxc )
{
maxc = cnt[c[x]];
sum = c[x];
}else if( cnt[c[x]] == maxc ) //两个颜色都为最大值都得算
{
sum += c[x];
}
for (int i = 0;i < g[x].size(); i++)
{
int t = g[x][i];
if( t == f || t == flag ) continue;
count(t,x,val);
}
}
void dfs2(int x,int f,int keep) //当前节点为x,父亲节点为f,keep为false表示不保留贡献
{
for (int i = 0; i < g[x].size(); i++) //处理轻儿子的子树,不保留贡献
{
int t = g[x][i];
if( t == f || t == son[x] ) continue;
dfs2(t,x,false);
}
if( son[x] )
{
dfs2(son[x],x,true); //处理重儿子的子树,保留贡献
flag = son[x]; //标记重儿子,这样在算贡献时不会算上它
}
count(x,f,1); //将轻儿子和自身的贡献加上
ans[x] = sum;
flag = 0;
if( !keep )
{
count(x,f,-1); //keep为false不保留,以x为子树的贡献都加上-1,即消去了贡献
sum = maxc = 0; //消去贡献后,sum和maxc都需要置为0
}//没消去贡献时,maxc与sum都是有用的
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(0);
int n;
cin >> n;
for (int i = 1; i <= n; i++)
{
cin >> c[i];
}
for (int i = 1; i < n; i++)
{
int x,y;
cin >> x >> y;
g[x].push_back(y);
g[y].push_back(x);
}
dfs1(1,0);
dfs2(1,0,false);
for (int i = 1; i <= n; i++)
{
cout << ans[i];
if( i == n ) cout << '\n';
else cout << ' ';
}
return 0;
}
证明复杂度:
对于一个节点需要再次被计算,当且仅当它的某个祖先节点不为其父亲节点的重儿子,假设这个节点被计算了logn次,由于有重儿子的存在,每次计算都会附带一个重儿子,节点数每次至少乘2,logn次后就达到了n,所以一个节点被计算次数不超过logn。
更多推荐
已为社区贡献2条内容
所有评论(0)