Codeforces600E Lomsat gelral(线段树合并)

题目描述

You are given a rooted tree with root in vertex 1 1 1 . Each vertex is coloured in some colour.

Let's call colour c c c dominating in the subtree of vertex v v v if there are no other colours that appear in the subtree of vertex v v v more times than colour c c c . So it's possible that two or more colours will be dominating in the subtree of some vertex.

The subtree of vertex v v v is the vertex v v v and all other vertices that contains vertex v v v in each path to the root.

For each vertex v v v find the sum of all dominating colours in the subtree of vertex v v v .

输入输出格式

输入格式：

The first line contains integer n n n ( 1<=n<=105 1<=n<=10^{5} 1<=n<=105 ) — the number of vertices in the tree.

The second line contains n n n integers ci c_{i} ci​ ( 1<=ci<=n 1<=c_{i}<=n 1<=ci​<=n ), ci c_{i} ci​ — the colour of the i i i -th vertex.

Each of the next n−1 n-1 n−1 lines contains two integers xj,yj x_{j},y_{j} xj​,yj​ ( 1<=xj,yj<=n 1<=x_{j},y_{j}<=n 1<=xj​,yj​<=n ) — the edge of the tree. The first vertex is the root of the tree.

输出格式：

Print n n n integers — the sums of dominating colours for each vertex.

输入输出样例

输入样例#1： 复制

4
1 2 3 4
1 2
2 3
2 4

输出样例#1： 复制

10 9 3 4

输入样例#2： 复制

15
1 2 3 1 2 3 3 1 1 3 2 2 1 2 3
1 2
1 3
1 4
1 14
1 15
2 5
2 6
2 7
3 8
3 9
3 10
4 11
4 12
4 13

输出样例#2： 复制

6 5 4 3 2 3 3 1 1 3 2 2 1 2 3

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

题解：就是把每个点的颜色先插入每个点的对应的线段树里，然后dfs时把子树的线段树合并到父节点上，维护题意要求的值就可以了

代码如下：

#include<cstdio>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
#define lson tr[now].l
#define rson tr[now].r
#define int long long
using namespace std;

struct tree
{
    int l,r,sum,val,ans;
}tr[5000050];

int rt[100010],cl[100010],cnt,n,anss[100010];
vector<int> g[100010];

int push_up(int now)
{
    if(tr[lson].sum>tr[rson].sum)
    {
        tr[now].sum=tr[lson].sum;
        tr[now].val=tr[lson].val;
        tr[now].ans=tr[lson].ans;
    }
    if(tr[rson].sum>tr[lson].sum)
    {
        tr[now].sum=tr[rson].sum;
        tr[now].val=tr[rson].val;
        tr[now].ans=tr[rson].ans;
    }
    if(tr[lson].sum==tr[rson].sum)
    {
        tr[now].sum=tr[lson].sum;
        tr[now].val=tr[lson].val;
        tr[now].ans=tr[lson].ans+tr[rson].ans;
    }
}

int update(int &now,int l,int r,int pos,int v)
{
    if(!now) now=++cnt;
    if(l==r)
    {
        tr[now].val=l;
        tr[now].sum+=v;
        tr[now].ans=l;
        return 0;
    }
    int mid=(l+r)>>1;
    if(pos<=mid) update(lson,l,mid,pos,v);
    else update(rson,mid+1,r,pos,v);
    push_up(now);
}

int merge(int a,int b,int l,int r)
{
    if(!a) return b;
    if(!b) return a;
    if(l==r)
    {
        tr[a].val=l;
        tr[a].sum+=tr[b].sum;
        tr[a].ans=l;
        return a;
    }
    int mid=(l+r)>>1;
    tr[a].l=merge(tr[a].l,tr[b].l,l,mid);
    tr[a].r=merge(tr[a].r,tr[b].r,mid+1,r);
    push_up(a);
    return a;
}

int dfs(int now,int f)
{
    for(int i=0;i<g[now].size();i++)
    {
        if(g[now][i]==f) continue;
        dfs(g[now][i],now);
        merge(rt[now],rt[g[now][i]],1,100000);
    }
    update(rt[now],1,100000,cl[now],1);
    anss[now]=tr[rt[now]].ans;
}

int main()
{
    scanf("%lld",&n);
    for(int i=1;i<=n;i++)
    {
        scanf("%lld",&cl[i]);
        rt[i]=i;
        cnt++;
    }
    int from,to;
    for(int i=1;i<n;i++)
    {
        scanf("%lld%lld",&from,&to);
        g[from].push_back(to);
        g[to].push_back(from);
    }
    dfs(1,0);
    for(int i=1;i<=n;i++)
    {
        printf("%lld ",anss[i]);
    }
}