[CF600E]Lomsat gelral

Description:

一棵树有n个结点，每个结点都是一种颜色，每个颜色有一个编号，求树中每个子树的最多的颜色编号的和。

Hint:

\(n \le 10^5\)

Solution:

线段树合并裸题?

#include <map>
#include <set>
#include <stack>
#include <cmath>
#include <queue>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <iostream>
#include <algorithm>
using namespace std;
typedef long long ll;
const int mxn=5e5+5;
int n,m,s,cnt,a[mxn],rt[mxn<<4],ls[mxn<<4],rs[mxn<<4],tot[mxn<<4];
ll sum[mxn<<4],ans[mxn];
int hd[mxn];

inline int read() {
	char c=getchar(); int x=0,f=1;
	while(c>'9'||c<'0') {if(c=='-') f=-1;c=getchar();}
	while(c<='9'&&c>='0') {x=(x<<3)+(x<<1)+(c&15);c=getchar();}
	return x*f;
}
inline void chkmax(int &x,int y) {if(x<y) x=y;}
inline void chkmin(int &x,int y) {if(x>y) x=y;}

struct ed {
	int to,nxt;
}t[mxn<<1];

inline void add(int u,int v) {
	t[++cnt]=(ed) {v,hd[u]}; hd[u]=cnt;
}

int merge(int u,int v,int l,int r) {
	if(!u) return v;
	if(!v) return u;
	if(l==r) {
		tot[u]+=tot[v];
		return u;
	}
	int mid=(l+r)>>1;
	ls[u]=merge(ls[u],ls[v],l,mid);
	rs[u]=merge(rs[u],rs[v],mid+1,r);
	tot[u]=max(tot[ls[u]],tot[rs[u]]); ll res=0;
	if(tot[u]==tot[rs[u]]) res+=sum[rs[u]];
	if(tot[u]==tot[ls[u]]) res+=sum[ls[u]];
	sum[u]=res;
	return u;
}

void update(int l,int r,int pos,int val,int &p) {
	if(!p) p=++s;
	if(l==r) {++tot[p]; sum[p]=l; return ;}
	int mid=(l+r)>>1;
	if(pos<=mid) update(l,mid,pos,val,ls[p]);
	else update(mid+1,r,pos,val,rs[p]);
    tot[p]=max(tot[ls[p]],tot[rs[p]]); ll res=0;
    if(tot[p]==tot[rs[p]]) res+=sum[rs[p]];
    if(tot[p]==tot[ls[p]]) res+=sum[ls[p]];
	sum[p]=res;
}

void dfs(int u,int fa) {
	for(int i=hd[u];i;i=t[i].nxt) {
		int v=t[i].to;
		if(v==fa) continue ;
		dfs(v,u); rt[u]=merge(rt[u],rt[v],1,n);
	}
	update(1,n,a[u],1,rt[u]);
	ans[u]=sum[rt[u]];
}

int main()
{
	n=read(); int u,v;
	for(int i=1;i<=n;++i) a[i]=read();
	for(int i=1;i<n;++i) {
		u=read(); v=read();
		add(u,v); add(v,u);
	}
	dfs(1,0);
	for(int i=1;i<=n;++i) printf("%I64d\n",ans[i]);
    return 0;
}