树上颜色交 题解

__ryp__ · 2024-09-25 07:30:34 · 个人记录

考虑到两个子树的并就是整个树，那么一个颜色在两个子树同时存在，当且仅当它在某一个子树中的出现次数不为它在整棵树里头的出现次数。

那么就做完了。因为这时候我们就可以钦定 1 为根，然后对 1 为根的每一个子树算一下它的答案。线段树维护这个是简单的。

然后可以线段树合并 / 树上启发式合并，做完了。

#include <iostream>
#include <vector>
using namespace std;
const int N = 1e5 + 10;
int rt[N], z[N], cntq[N], head[N], res[N], ecnt = 1, cnt = 0, n;
struct edge { int v, nxt; } e[N * 2];
struct node { int s, k, l, r; } sg[50 * N];
void add (int u, int v)
{
    e[++ecnt] = { v, head[u] }, head[u] = ecnt;
}

// 考虑到实际上就是不等于总出现次数的点的数量。可以直接线段树做。

void push_up (int i) { sg[i].s = sg[sg[i].l].s + sg[sg[i].r].s; }
void upd (int &u, int x, int y, int p)
{
    int mid = (x + y) / 2;
    if (!u) u = ++cnt;
    if (x == y) return sg[u].s = cntq[x] != ++sg[u].k, void ();
    if (p <= mid) upd (sg[u].l, x, mid, p);
    else upd (sg[u].r, mid + 1, y, p);
    push_up (u);
}

void merge (int &u, int v, int x, int y)
{
    int mid = (x + y) / 2;
    if (!u || !v) return u += v, void ();
    if (x == y) return sg[u].s = (sg[u].k += sg[v].k) != cntq[x], void ();
    merge (sg[u].l, sg[v].l, x, mid), merge (sg[u].r, sg[v].r, mid + 1, y);
    push_up (u);
}

void dfs (int u, int last)
{
    int nr = 0;
    for (int i = head[u]; i; i = e[i].nxt) {
        int v = e[i].v;
        if (v == last) nr = i / 2;
        else dfs (v, u), merge (rt[u], rt[v], 1, n);
    }
    upd (rt[u], 1, n, z[u]), res[nr] = sg[rt[u]].s;
}

int main (void)
{
    cin.tie (0)->sync_with_stdio (false);
    cin >> n;
    for (int i = 1; i <= n; i++) cin >> z[i], cntq[z[i]]++;
    for (int i = 1, u, v; i < n; i++) cin >> u >> v, add (u, v), add (v, u);
    dfs (1, 0);
    for (int i = 1; i < n; i++) cout << res[i] << '\n';
    return 0;
}