codeforces 600E E. Lomsat gelral (线段树合并)

codeforces 600E E. Lomsat gelral

传送门:https://codeforces.com/contest/600/problem/E

题意:

给你一颗n个节点的树,树上的每一个节点都有一种颜色,询问每一个节点所在的子树颜色数量最多的那些颜色的值的总和

题解:

维护子树颜色的数量和答案,线段树合并即可

代码:

/**
 *        ┏┓    ┏┓
 *        ┏┛┗━━━━━━━┛┗━━━┓
 *        ┃       ┃  
 *        ┃   ━    ┃
 *        ┃ >   < ┃
 *        ┃       ┃
 *        ┃... ⌒ ...  ┃
 *        ┃       ┃
 *        ┗━┓   ┏━┛
 *          ┃   ┃ Code is far away from bug with the animal protecting          
 *          ┃   ┃   神兽保佑,代码无bug
 *          ┃   ┃           
 *          ┃   ┃        
 *          ┃   ┃
 *          ┃   ┃           
 *          ┃   ┗━━━┓
 *          ┃       ┣┓
 *          ┃       ┏┛
 *          ┗┓┓┏━┳┓┏┛
 *           ┃┫┫ ┃┫┫
 *           ┗┻┛ ┗┻┛
 */
// warm heart, wagging tail,and a smile just for you!
//
//                            _ooOoo_
//                           o8888888o
//                           88" . "88
//                           (| -_- |)
//                           O\  =  /O
//                        ____/`---'\____
//                      .'  \|     |//  `.
//                     /  \|||  :  |||//  \
//                    /  _||||| -:- |||||-  \
//                    |   | \  -  /// |   |
//                    | \_|  ''\---/''  |   |
//                    \  .-\__  `-`  ___/-. /
//                  ___`. .'  /--.--\  `. . __
//               ."" '<  `.___\_<|>_/___.'  >'"".
//              | | :  `- \`.;`\ _ /`;.`/ - ` : | |
//              \  \ `-.   \_ __\ /__ _/   .-` /  /
//         ======`-.____`-.___\_____/___.-`____.-'======
//                            `=---='
//        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
//                     佛祖保佑      永无BUG
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
using namespace std;
typedef long long LL;
typedef pair pii;
typedef unsigned long long uLL;
#define bug printf("*********\n")
#define FIN freopen("input.txt","r",stdin);
#define FON freopen("output.txt","w+",stdout);
#define IO ios::sync_with_stdio(false),cin.tie(0)
#define debug1(x) cout<<"["<<#x<<" "<<(x)<<"]\n"
#define debug2(x,y) cout<<"["<<#x<<" "<<(x)<<" "<<#y<<" "<<(y)<<"]\n"
#define debug3(x,y,z) cout<<"["<<#x<<" "<<(x)<<" "<<#y<<" "<<(y)<<" "<<#z<<" "<>= 1;
    } return ans;
}
struct node {
    int l, r, sum, val;
    LL ans;
} tree[maxn * 50];
int root[maxn];
int col[maxn];
int tree_cnt;
struct EDGE {
    int v, nxt;
} edge[maxn << 1];
int head[maxn], tot;

void add_edge(int u, int v) {
    edge[tot].v = v;
    edge[tot].nxt = head[u];
    head[u] = tot++;
}
#define ls tree[rt].l
#define rs tree[rt].r
void push_up(int rt) {
    if(tree[ls].sum == tree[rs].sum) {
        tree[rt].sum = tree[ls].sum;
        tree[rt].val = tree[ls].sum;
        tree[rt].ans = tree[ls].ans + tree[rs].ans;
    } else if(tree[ls].sum > tree[rs].sum) {
        tree[rt].sum = tree[ls].sum;
        tree[rt].val = tree[ls].val;
        tree[rt].ans = tree[ls].ans;
    } else {
        tree[rt].sum = tree[rs].sum;
        tree[rt].val = tree[rs].val;
        tree[rt].ans = tree[rs].ans;
    }
}
void update(int &x, int l, int r, int pos, int val) {
    if(!x) x = ++tree_cnt;
    if(l == r) {
        tree[x].val = l;
        tree[x].sum += val;
        tree[x].ans = l;
        return;
    }
    int mid = (l + r) >> 1;
    if(pos <= mid) update(tree[x].l, l, mid, pos, val);
    else update(tree[x].r, mid + 1, r, pos, val);
    push_up(x);
}
int merge(int x, int y, int l, int r) {
    if(!x) return y;
    if(!y) return x;
    if(l == r) {
        tree[x].val = l;
        tree[x].sum += tree[y].sum;
        tree[x].ans = l;
        return x;
    }
    int mid = (l + r) >> 1;
    tree[x].l = merge(tree[x].l, tree[y].l, l, mid);
    tree[x].r = merge(tree[x].r, tree[y].r, mid + 1, r);
    push_up(x);
    return x;
}
int Max = 100000;
LL ans[maxn];
void dfs(int u, int fa) {
    for(int i = head[u]; i != -1; i = edge[i].nxt) {
        int v = edge[i].v;
        if(v == fa) continue;
        dfs(v, u);
        merge(root[u], root[v], 1, Max);
    }
    update(root[u], 1, Max, col[u], 1);
    ans[u] = tree[root[u]].ans;
}
void init() {
    memset(head, -1, sizeof(head));
    memset(col, 0, sizeof(col));
    tree_cnt = tot = 0;

}
int main() {
#ifndef ONLINE_JUDGE
    FIN
#endif
    int n;
    scanf("%d", &n);
    init();
    for(int i = 1; i <= n; i++) {
        scanf("%d", &col[i]);
        root[i] = i;
        tree_cnt++;
    }
    for(int i = 1, u, v; i < n; i++) {
        scanf("%d%d", &u, &v);
        add_edge(u, v);
        add_edge(v, u);
    }
    dfs(1, 0);
    for(int i = 1; i <= n; i++) {
        printf("%lld%c", ans[i], i == n ? '\n' : ' ');
    }
    return 0;
}

你可能感兴趣的:(codeforces 600E E. Lomsat gelral (线段树合并))