Educational Codeforces Round 2 E. Lomsat gelral (权值线段树,动态开点+线段树启发式合并)

Educational Codeforces Round 2 E. Lomsat gelral (权值线段树,动态开点+线段树启发式合并)

题目

600 E. Lomsat gelral

题意

给你一个树（根节点为1），有n个节点，每个节点都有一种颜色，问你每个节点和其子节点一共有多少种颜色。(n<=1e5)

题解

最朴素的想法肯定是疯狂dfs，显然这是O(n²)的做法，不可取。
那么我们就要考虑到用树上启发式合并，这里我用的是权值线段树启发式合并。
思路是这样的：
从根节点开始dfs，为每一个叶子节点动态开点，创建一个线段树，每个节点深度做多是$logn$（所以我们线段树要开$nlogn$的空间，我开的是20*maxn）
然后再给每个父节点进行merge操作，merge操作具体是这样的：
先对每个父节点遍历到叶子节点，然后左右叶子节点之间互相合并。
这里要注意创建线段树和merge操作都需要pushup，把子节点的数据传给父节点。
这样代码就完成了。放一下AC代码

代码

#include <bits/stdc++.h>
#define pb push_back
using namespace std;
const int maxn=1e5+10;
typedef long long ll;

int n,cor[maxn];
ll ans[maxn];
vector<int> G[maxn];
struct node
{
#define ls(x) t[x].lc
#define rs(x) t[x].rc
	int lc,rc,cnt;
	ll sum;
} t[maxn*20];//线段树空间要开足够大
int tot,rt[maxn];
inline void pushup(int x)
{
	if(t[ls(x)].cnt == t[rs(x)].cnt)//左右子树的数量相同 
	{
		t[x].cnt=t[ls(x)].cnt;
		t[x].sum=t[ls(x)].sum+t[rs(x)].sum;
	}
	else if(t[ls(x)].cnt > t[rs(x)].cnt)//左子树大于右子树 
	{
		t[x].cnt=t[ls(x)].cnt;
		t[x].sum=t[ls(x)].sum;
	}
	else 
	{
		t[x].cnt=t[rs(x)].cnt;
		t[x].sum=t[rs(x)].sum;
	}
}
void insert(int &x,int l,int r,int pos)//权值线段树,动态开点
{
	if(!x)x=++tot;
	if(l==r)
	{
		t[x].sum=l;//权值为l 
		t[x].cnt++;//数量+1 
		return;
	}
	int mid=l+r>>1;
	if(pos<=mid)
		insert(ls(x),l,mid,pos);
	else
		insert(rs(x),mid+1,r,pos);
	pushup(x);
}
int merge(int x,int y,int l,int r)
{
	if(!x||!y)//如果有一个根节点是0,直接返回x+y
		return x+y;
	if(l==r)//如果是同一个节点，则重合 
	{
		t[x].sum=l;
		t[x].cnt+=t[y].cnt;
		return x;
	}
	int mid=l+r>>1; 
	ls(x)=merge(ls(x),ls(y),l,mid);//左儿子合并 
	rs(x)=merge(rs(x),rs(y),mid+1,r);//右儿子合并 
	pushup(x);//更新线段树 
	return x;
}
void dfs(int u,int fa)
{
	for(auto v:G[u])
	{
		if(v==fa)continue;
		dfs(v,u);
		rt[u]=merge(rt[u],rt[v],1,n);
	}
	//遍历到最根节点创建线段树 
	insert(rt[u],1,n,cor[u]);
	ans[u]=t[rt[u]].sum;
}
int main()
{
	scanf("%d",&n);
	for(int i=1; i<=n; i++)scanf("%d",&cor[i]);
	for(int i=1; i<n; i++)
	{
		int x,y;
		scanf("%d%d",&x,&y);
		G[x].pb(y),G[y].pb(x);
	}
	dfs(1,0);
	for(int i=1; i<=n; i++)printf("%lld ",ans[i]);
	return 0;
}