2016湖南省赛 I Tree Intersection（线段树合并，树链剖分）

传送门：https://ac.nowcoder.com/acm/contest/1112/I

题意：

给你一个n个结点的树，树上每个节点有自己的颜色

问你删除第i条边后形成的两颗子树有多少个相同的颜色

题解：

树链剖分写法：

对于每一种颜色来说，如果这个颜色是在单独的一颗子树中，那么就不会对其他的边产生贡献，所以我们单独对每一种颜色对边的贡献讨论，如果这个颜色只有一个，那么就不会产生贡献，否则，他就可以在两个相同颜色之间的路径上的边产生贡献，用树剖剖下来后区间加一加就可以了

线段树合并的写法

对于每一种颜色建权值线段树，线段树合并统计答案即可

代码：

树链剖分写法：

/**

 *　　　　　　　　┏┓　　 　┏┓

 * 　　　　　　　┏┛┗━━━━━━━┛┗━━━┓

 * 　　　　　　　┃　　　　　　　┃ 　

 * 　　　　　　　┃　　　━　　 　┃

 * 　　　　　　　┃　＞　　　＜　┃

 * 　　　　　　　┃　　　　　　　┃

 * 　　　　　　　┃...　⌒　... 　┃

 * 　　　　　　　┃　　　　　　　┃

 * 　　　　　　　┗━┓　　　┏━┛

 * 　　　　　　　　　┃　　　┃　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 <set>

#include <map>

#include <stack>

#include <cmath>

#include <queue>

#include <cstdio>

#include <string>

#include <vector>

#include <cstring>

#include <iostream>

#include <algorithm>

using namespace std;

typedef long long LL;

typedef pair<int, int> pii;

typedef unsigned long long uLL;

#define ls rt<<1

#define rs rt<<1|1

#define lson l,mid,rt<<1

#define rson mid+1,r,rt<<1|1

#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<<" "<<z<<"]\n"

const int maxn = 3e5 + 5;

const int INF = 0x3f3f3f3f;

const int mod = 1e9 + 7;

const double Pi = acos(-1);

LL gcd(LL a, LL b) {

    return b ? gcd(b, a % b) : a;

}

LL lcm(LL a, LL b) {

    return a / gcd(a, b) * b;

}

double dpow(double a, LL b) {

    double ans = 1.0;

    while(b) {

        if(b % 2)ans = ans * a;

        a = a * a;

        b /= 2;

    } return ans;

}

LL quick_pow(LL x, LL y) {

    LL ans = 1;

    while(y) {

        if(y & 1) {

            ans = ans * x % mod;

        } x = x * x % mod;

        y >>= 1;

    } return ans;

}

int n;

int a[maxn], sz[maxn], dep[maxn], fa[maxn], top[maxn], w[maxn], son[maxn], W[maxn];

int sum[maxn << 2], lazy[maxn << 2], vis[maxn << 2];

struct EDGE {

    int u, v, nt;

} edge[maxn << 1];

int head[maxn], summ, cnt;

int u[maxn];

int v[maxn];

void add_edge(int u, int v) {

    edge[++summ].u = u;

    edge[summ].v = v;

    edge[summ].nt = head[u];

    head[u] = summ;

}

void dfs1(int u) {

    sz[u] = 1; son[u] = 0;

    for (int i = head[u]; ~i; i = edge[i].nt) {

        int v = edge[i].v;

        if (v != fa[u]) {

            fa[v] = u;

            dep[v] = dep[u] + 1;

            dfs1(v);

            sz[u] += sz[v];

            if (sz[v] > sz[son[u]]) son[u] = v;

        }

    }

}

void dfs2(int u, int tp, int x) {

    top[u] = tp; w[u] = ++cnt; W[cnt] = u;

    if (son[u]) dfs2(son[u], tp, 1);

    for (int i = head[u]; ~i; i = edge[i].nt) {

        int v = edge[i].v;

        if (v == son[u] || v == fa[u]) continue;

        dfs2(v, v, 2);

    }

}

vector<int> vec[maxn];

void init() {

    memset(head, -1, sizeof(head));

    summ = 1; cnt = 0;

    for(int i = 1; i <= n; i++) {

        vec[i].clear();

    }

    for (int i = 1; i <= n; i++) {

        scanf("%d", &a[i]);

        vec[a[i]].push_back(i);

    }

    for (int i = 1; i < n; i++) {

        // int u, v;

        scanf("%d %d", &u[i], &v[i]);

        add_edge(u[i], v[i]);

        add_edge(v[i], u[i]);

    }

    dep[1] = 1; fa[1] = 0;

    dfs1(1);

    dfs2(1, 1, 1);

}

void pushup(int rt) {

    if(vis[ls] == vis[rs]) vis[rt] = vis[ls];

    else vis[rt] = -1;

}

void pushdown(int rt, int mid) {

    if (lazy[rt]) {

        lazy[rt << 1] = lazy[rt];

        lazy[rt << 1 | 1] = lazy[rt];

        vis[rt << 1] = 0;

        vis[rt << 1 | 1] = 0;

        lazy[rt] = 0;

    }

}

void build(int l, int r, int rt) {

    lazy[rt] = 0;

    sum[rt] = 0;

    if (l == r) {

        vis[rt] = 0;

        return;

    }

    int mid = (l + r) >> 1;

    build(lson);

    build(rson);

}

void update(int L, int R, int val, int l, int r, int rt) {

    // debug2(L, R);

    if(L > R || r < L || R < l) return;

    if(vis[rt] == 1) return;

    if (L <= l && r <= R) {

        // sum[rt] += val * (r - l + 1);

        // lazy[rt] += val;

        if(vis[rt] == 0) {

            vis[rt] = 1;

            sum[rt]++;

            return;

        }

    }

    pushdown(rt, r - l + 1);

    int mid = (l + r) >> 1;

    if (L <= mid) update(L, R, val, lson);

    if (R > mid) update(L, R, val, rson);

    pushup(rt);

}

int query(int pos, int l, int r, int rt) {

    if (l == r) {

        return sum[rt];

    }

    pushdown(rt, r - l + 1);

    int mid = (l + r) >> 1;

    int ans = sum[rt];

    if (pos <= mid)  ans += query(pos, lson);

    else  ans += query(pos, rson);

    return ans;

}

int LCA(int x, int y) {

    // debug2(x, y);

    while (top[x] != top[y]) {

        // bug;

        if (dep[top[x]] < dep[top[y]]) swap(x, y);

        update(w[top[x]], w[x], 1, 1, n, 1);

        x = fa[top[x]];

    }

    if (dep[x] > dep[y])

        std::swap(x, y);

    update(w[x] + 1, w[y], 1, 1, n, 1);

    return x;

}

int main() {

#ifndef ONLINE_JUDGE

    FIN

#endif

    while(scanf("%d", &n) != EOF) {

        init();

        build(1, n, 1);

        // for(int i = 1; i <= n; i++) {

        //     // printf("%d\n", w[]);

        //     // debug1(w[i]);

        // }

        for(int i = 1; i <= n; i++) {

            vis[1] = 0;

            lazy[1] = 1;

            int sz = vec[i].size();

            if(sz < 2) continue;

            int bb = LCA(vec[i][0], vec[i][1]);

            for(int j = 2; j < sz; j++) {

                // debug2(vec[i][0], vec[i][1]);

                bb = LCA(bb, vec[i][j]);

            }

        }

        for(int i = 1; i < n; i++) {

            int ans;

            // debug2(v[i], u[i]);

            if(fa[v[i]] == u[i]) {

                ans = query(w[v[i]], 1, n, 1);

            } else {

                ans = query(w[u[i]], 1, n, 1);

            }

            printf("%d\n", ans);

        }

    }

    return 0;

}

线段树合并写法

/**

 *　　　　　　　　┏┓　　 　┏┓

 * 　　　　　　　┏┛┗━━━━━━━┛┗━━━┓

 * 　　　　　　　┃　　　　　　　┃ 　

 * 　　　　　　　┃　　　━　　 　┃

 * 　　　　　　　┃　＞　　　＜　┃

 * 　　　　　　　┃　　　　　　　┃

 * 　　　　　　　┃...　⌒　... 　┃

 * 　　　　　　　┃　　　　　　　┃

 * 　　　　　　　┗━┓　　　┏━┛

 * 　　　　　　　　　┃　　　┃　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 <set>

#include <map>

#include <stack>

#include <cmath>

#include <queue>

#include <cstdio>

#include <string>

#include <bitset>

#include <vector>

#include <cstring>

#include <iostream>

#include <algorithm>

using namespace std;

typedef long long LL;

typedef pair<int, int> pii;

typedef unsigned long long uLL;

#define ls rt<<1

#define rs rt<<1|1

#define lson l,mid,rt<<1

#define rson mid+1,r,rt<<1|1

#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<<" "<<z<<"]\n"

const int maxn = 1e5 + 5;

const int INF = 0x3f3f3f3f;

const int mod = 1e9 + 7;

const double Pi = acos(-1);

LL gcd(LL a, LL b) {

    return b ? gcd(b, a % b) : a;

}

LL lcm(LL a, LL b) {

    return a / gcd(a, b) * b;

}

double dpow(double a, LL b) {

    double ans = 1.0;

    while(b) {

        if(b % 2)ans = ans * a;

        a = a * a;

        b /= 2;

    } return ans;

}

LL quick_pow(LL x, LL y) {

    LL ans = 1;

    while(y) {

        if(y & 1) {

            ans = ans * x % mod;

        } x = x * x % mod;

        y >>= 1;

    } return ans;

}

int col[maxn];

int cnt[maxn];

struct EDGE {

    int v, nxt;

} edge[maxn << 1];

int head[maxn], tot;

void init() {

    memset(head, -1, sizeof(head));

    tot = 0;

}

void add_edge(int u, int v) {

    edge[tot].v = v;

    edge[tot].nxt = head[u];

    head[u] = tot++;

}

struct node {

    int l, r, sum, val;

} tree[maxn * 20];

int root[maxn];

int rt_cnt = 0;

int n;

int ans[maxn];

void push_up(int rt) {

    tree[rt].sum = tree[tree[rt].l].sum + tree[tree[rt].r].sum;

}

int build(int l, int r, int pos) {

    int rt = ++rt_cnt;

    tree[rt].l = 0;

    tree[rt].r = 0;

    tree[rt].sum = 0;

    if(l == r) {

        tree[rt].val = 1;

        if(tree[rt].val != cnt[l]) {

            tree[rt].sum = 1;

        } else {

            tree[rt].sum = 0;

        }

        return rt;

    }

    int mid = (l + r) >> 1;

    if(pos <= mid) tree[rt].l = build(l, mid, pos);

    else tree[rt].r = build(mid + 1, r, pos);

    push_up(rt);

    return rt;

}

void merge(int &x, int y, int l, int r) {

    if(!x || !y) {

        if(!x) x = y;

        return;

    }

    if(l == r) {

        tree[x].val += tree[y].val;

        tree[x].sum = (tree[x].val != cnt[l]);

        return;

    }

    int mid = (l + r) >> 1;

    merge(tree[x].l, tree[y].l, l, mid);

    merge(tree[x].r, tree[y].r, mid + 1, r);

    push_up(x);

}

void dfs(int u, int fa, int id) {

    root[u] = build(1, n, col[u]);

    for(int i = head[u]; i != -1; i = edge[i].nxt) {

        int v = edge[i].v;

        if(v == fa) continue;

        dfs(v, u, i);

        merge(root[u], root[v], 1, n);

    }

    if(u != 1) {

        int tid = id / 2 + 1;

        ans[tid] = tree[root[u]].sum;

    }

}

int main() {

#ifndef ONLINE_JUDGE

    FIN

#endif

    scanf("%d", &n);

    for(int i = 1; i <= n; i++) {

        scanf("%d", &col[i]);

        // vec[col[i]].push_back(i);

        cnt[col[i]]++;

    }

    init();

    for(int i = 1, u, v; i < n; i++) {

        scanf("%d%d", &u, &v);

        add_edge(u, v);

        add_edge(v, u);

    }

    // debug1(n);

    dfs(1, -1, -1);

    for(int i = 1; i <= n - 1; i++) {

        printf("%d\n", ans[i]);

    }

    return 0;

}