Educational Codeforces Round 2 E. Lomsat gelral——线段树合并

题目大意

给你一棵树，每个树的结点都有一个颜色编号，问你对于每个结点，这个结点及其子节点出现过的颜色次数最多的编号之和是多少（题意有点绕）

解题思路

我们对于每一个结点建立一个权值线段树，树上的记录这个结点出现过最多次数的颜色编号之和（也就是ans），以及颜色最多出现的次数num。然后在dfs题目所给树的时候合并这些节点树并更新答案。
怎么pushup呢？
当左子树的num < 右子树num时，我们的ans和num都取右子树的
当左子树的num > 右子树num时，我们的ans和num都取左子树的
当左子树的num == 右子树num时，我们的ans取左右子树的ans之和（因为此时左右子树的最多颜色数量相同），num取左子树或者右子树
如何merge，当我们到达叶子节点的时候，我们就取两棵树的num之和，ans随便取其一就行，因为只有相同颜色的两棵树才能同时到叶子节点，所以他们的ans是相同的

Code

#include <bits/stdc++.h>
#define ll long long
#define qc ios::sync_with_stdio(false); cin.tie(0);cout.tie(0)
#define fi first
#define se second
#define PII pair<int, int>
#define PLL pair<ll, ll>
#define pb push_back
using namespace std;
const int MAXN = 2e6 + 7;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
const ll mod = 1e9 + 7;
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while (!isdigit(ch)){if (ch=='-') f=-1;ch=getchar();}
    while (isdigit(ch)){x=x*10+ch-48;ch=getchar();}
    return x*f;
}
int n;
ll a[MAXN], ans[MAXN];
struct Tree{
	int l, r, num;
	ll dat;
}tree[MAXN << 1];
int rt[MAXN];
vector<int> v[MAXN];
int tot;
int build(){
	tot++;
	tree[tot].l = tree[tot].r = tree[tot].num = tree[tot].dat = 0;
	return tot;
}
void pushup(int p){
	if(tree[tree[p].l].num > tree[tree[p].r].num){
		tree[p].num = tree[tree[p].l].num;
		tree[p].dat = tree[tree[p].l].dat;
	}
	else if(tree[tree[p].l].num < tree[tree[p].r].num){
		tree[p].num = tree[tree[p].r].num;
		tree[p].dat = tree[tree[p].r].dat;
	}
	else{
		tree[p].num = tree[tree[p].r].num;
		tree[p].dat = tree[tree[p].r].dat + tree[tree[p].l].dat;
	}
}
void update(int &p, int l, int r, int val, int pos){
	if(l == r){
		tree[p].num += val;
		tree[p].dat = pos;
		return ;
	}
	int m = (l + r) >> 1;
	if(pos <= m){
		if(!tree[p].l) tree[p].l = build();
		update(tree[p].l, l, m, val, pos);
	}
	else{
		if(!tree[p].r) tree[p].r = build();
		update(tree[p].r, m+1, r, val, pos);
	}
	pushup(p);
}
int merge(int p, int q, int l, int r){
	if(!p) return q;
	if(!q) return p;
	if(l == r){
		tree[p].num += tree[q].num;
		tree[p].dat = l;
		return p;
	}
	int m = (l + r) >> 1;
	tree[p].l = merge(tree[p].l, tree[q].l, l, m);
	tree[p].r = merge(tree[p].r, tree[q].r, m+1, r);
	pushup(p);
	return p;
}
void dfs(int x, int f){
	update(x, 1, n, 1, a[x]);
	for(auto i : v[x]){
		if(i == f)
			continue;
		dfs(i, x);
		merge(x, i, 1, n);
	}
	ans[x] = tree[x].dat;
}
void solve(){
    cin >> n;
    for(int i = 1; i <= n; i++){
    	cin >> a[i];
    	rt[i] = i;
    	tot++;
    }
    for(int i = 1; i < n; i++){
    	int x, y;
    	cin >> x >> y;
    	v[x].pb(y);
    	v[y].pb(x);
    }
    dfs(1, 0);
    for(int i = 1; i <= n; i++)
    	cout << ans[i] << " ";
    cout << endl;
}

int main()
{
	#ifdef ONLINE_JUDGE
    #else
       freopen("in.txt", "r", stdin);
       freopen("out.txt", "w", stdout);
    #endif

	qc;
    int T;
    // cin >> T;
    T = 1;
    while(T--){

        solve();
    }
	return 0;
}