CSU1811: Tree Intersection

题目大意

给一棵点带颜色的树，问依次删掉每条边后得到的两棵子树出现的颜色并集的大小。

简要题解

注意这么一个性质，一种颜色在两棵树中出现，则在一棵树中出现，且出现次数小于总的出现次数。

然后需要这么一个数据结构，能维护子树内出现的颜色和对应的出现次数。

用数组启发式合并或者线段树合并即可。

之前写了一发，拍了好久没问题，交上去就RE，好气。

今天重新写了一发，写完直接AC，神清气爽hhh

#include <bits/stdc++.h>
using namespace std;
namespace my_header {
#define pb push_back
#define mp make_pair
#define pir pair<int, int>
#define vec vector<int>
#define pc putchar
#define clr(t) memset(t, 0, sizeof t)
#define pse(t, v) memset(t, v, sizeof t)
#define bl puts("")
#define wn(x) wr(x), bl
#define ws(x) wr(x), pc(' ')
    const int INF = 0x3f3f3f3f;
    typedef long long LL;
    typedef double DB;
    inline char gchar() {
        char ret = getchar();
        for(; (ret == '
' || ret == '
' || ret == ' ') && ret != EOF; ret = getchar());
        return ret; }
    template<class T> inline void fr(T &ret, char c = ' ', int flg = 1) {
        for(c = getchar(); (c < '0' || '9' < c) && c != '-'; c = getchar());
        if (c == '-') { flg = -1; c = getchar(); }
        for(ret = 0; '0' <= c && c <= '9'; c = getchar())
            ret = ret * 10 + c - '0';
        ret = ret * flg; }
    inline int fr() { int t; fr(t); return t; }
    template<class T> inline void fr(T&a, T&b) { fr(a), fr(b); }
    template<class T> inline void fr(T&a, T&b, T&c) { fr(a), fr(b), fr(c); }
    template<class T> inline char wr(T a, int b = 10, bool p = 1) {
        return a < 0 ? pc('-'), wr(-a, b, 0) : (a == 0 ? (p ? pc('0') : p) : 
            (wr(a/b, b, 0), pc('0' + a % b)));
    }
    template<class T> inline void wt(T a) { wn(a); }
    template<class T> inline void wt(T a, T b) { ws(a), wn(b); }
    template<class T> inline void wt(T a, T b, T c) { ws(a), ws(b), wn(c); }
    template<class T> inline void wt(T a, T b, T c, T d) { ws(a), ws(b), ws(c), wn(d); }
    template<class T> inline T gcd(T a, T b) {
        return b == 0 ? a : gcd(b, a % b); }
    template<class T> inline T fpw(T b, T i, T _m, T r = 1) {
        for(; i; i >>= 1, b = b * b % _m)
            if(i & 1) r = r * b % _m;
        return r; }
};
using namespace my_header;

const int MAXN = 5e6 + 100;
const int MAXV = 2e5 + 100;
int n, col[MAXV], fa[MAXV], cnt[MAXV], ans[MAXV];
int h[MAXV], nxt[MAXV], to[MAXV], e;

#define lc ch[0]
#define rc ch[1]
struct SegNode {
    SegNode *ch[2];
    int s, c;
    void init() {
        s = c = 0;
        lc = rc = NULL;
    }
    void update();
    void merge(int, int, SegNode*);
    void insert(int, int, int);
} node_pool[MAXN], *loc, *root[MAXV];
SegNode* newSegNode() {
    loc->init();
    return loc++;
}

void SegNode::update() {
    s = (lc ? lc->s : 0) + (rc ? rc->s : 0);
}
void SegNode::insert(int l, int r, int v) {
    if (l == r) {
        ++c;
        s = (c != cnt[l]);
    } else {
        int m = (l + r) >> 1;
        if (v <= m) {
            if (!lc)
                lc = newSegNode();
            lc->insert(l, m, v);
        } else {
            if (!rc)
                rc = newSegNode();
            rc->insert(m + 1, r, v);
        }
        update();
    }
}
void SegNode::merge(int l, int r, SegNode* v) {
    if (l == r) {
        if (v) {
            c += v->c;
            s = (c != cnt[l]);
        }
    } else {
        int m = (l + r) >> 1;
        if (!lc)
            lc = v->lc;
        else if (v->lc)
            lc->merge(l, m, v->lc);
        if (!rc)
            rc = v->rc;
        else if (v->rc)
            rc->merge(m + 1, r, v->rc);
        update();
    }
}

void dfs(int u) {
    int w = -1;
    root[u] = newSegNode();
    root[u]->insert(1, n, col[u]);
    for (int i = h[u]; i != -1; i = nxt[i]) {
        int v = to[i];
        if (v != fa[u]) {
            fa[v] = u;
            dfs(v);
            root[u]->merge(1, n, root[v]);
        } else w = i >> 1;
    }
    if (w != -1)
        ans[w] = root[u]->s;
}

int main() {
#ifdef lol
    freopen("1811.in", "r", stdin);
    freopen("1811.out", "w", stdout);
#endif

    while (scanf("%d", &n) != EOF) {
        for (int i = 1; i <= n; ++i)
            cnt[i] = 0;
        for (int i = 1; i <= n; ++i) {
            fr(col[i]);
            ++cnt[col[i]];
        }
        e = 0;
        memset(h, -1, sizeof h);
        for (int i = 1; i < n; ++i) {
            int u, v;
            fr(u, v);
            to[e] = v, nxt[e] = h[u], h[u] = e++;
            to[e] = u, nxt[e] = h[v], h[v] = e++;
        }
        loc = node_pool;
        dfs(1);
        for (int i = 0; i < n - 1; ++i)
            wt(ans[i]);
    }

    return 0;
}