题解 P5852 [USACO19DEC]Bessie's Snow Cow P

$\mathtt{\text{link}}$

题意

给出一颗 $n$ 个点的树，根为 $1$，定义点 $u$ 的权值为 $u$ 上不同的颜色的数量，$m$ 次操作：

1. 将 $u$ 的子树中没有颜色 $c$ 的点加上颜色 $c$。
2. 查询 $u$ 子树的权值和。

数据范围： $n,m,c\le {10}^5$。

题解

首先可以用 dfs 序把子树转移到区间上去。

发现一个点 $u$ 子树的权值和贡献有两种：

1. $u$ 祖先的子树均被染色，那点 $u$ 子树也会被染色，贡献就是子树大小。

2. $u$ 子树部分被染色，要加回来。

可以用两个 DS，一个维护每个点情况 $1$ 的贡献，区间加，单点查，一个维护该点情况 $2$ 对子树的贡献，单点加，区间查。直接树状数组维护。

发现两种情况会有重复部分，可对每个颜色开一个 set，维护情况 $1$ 中被染色的子树根节点的 dfn，新加入一个节点时，分两种情况：

1. 有祖先的子树被染色，此时不需要再加入。

2. 没有祖先的子树被染色，先把新节点子树中在 set 里的全部减掉，再把该节点加进去。转换到区间上就是把新的子树区间并上已经在 set 里的区间，把被包含的区间全部删掉即可。发现每次操作 set 最多加 $1$ 个节点，总数不会超过 $m$，复杂度有保证。

每次查询寻找 dfn 前驱复杂度 $O(\log n)$，树状数组维护是 $O(\log n)$，所以总复杂度为 $O(n+m\log n)$。

#include <bits/stdc++.h>
using namespace std;
namespace IO {
    //read and write
} using namespace IO;
const int N = 1e5 + 10;
int n, q;
struct bit {
    int t[N];
    void add(int x, int v) { for(int i = x; i <= n; i += i & -i) t[i] += v; }
    int qry(int x) { int res = 0; for(int i = x; i; i -= i & -i) res += t[i]; return res; }
} t1, t2;
struct edge {
    int to, nxt;
} e[N << 1];
int head[N], cnt;
void add(int u, int v) {
    e[++cnt] = (edge){v, head[u]};
    head[u] = cnt;
}
set<int> s[N];
set<int>::iterator it;
int dfn[N], siz[N], id[N], tim;
void dfs(int u, int fa) {
    dfn[u] = ++tim, id[tim] = u, siz[u] = 1;
    for(int i = head[u], v; i; i = e[i].nxt) {
        v = e[i].to;
        if(v == fa) continue;
        dfs(v, u);
        siz[u] += siz[v];
    }
}
int main() {
    n = read(), q = read();
    for(int i = 1, u, v; i < n; i++) {
        u = read(), v = read();
        add(u, v), add(v, u);
    }
    dfs(1, 0);
    for(int i = 1, opt, u, c; i <= q; i++) {
        opt = read(), u = read();
        if(opt == 1) {
            c = read();
            it = s[c].lower_bound(dfn[u]);
            if(it != s[c].begin() && dfn[id[*prev(it)]] + siz[id[*prev(it)]] >= dfn[u] + siz[u]) continue;
            while(it != s[c].end() && *it < dfn[u] + siz[u])
                t1.add(*it, -1), t1.add(*it + siz[id[*it]], 1), t2.add(*it, -siz[id[*it]]), s[c].erase(it++);
            s[c].insert(dfn[u]), t1.add(dfn[u], 1), t1.add(dfn[u] + siz[u], -1), t2.add(dfn[u], siz[u]);
        }
        else printf("%d\n", siz[u] * t1.qry(dfn[u]) + t2.qry(dfn[u] + siz[u] - 1) - t2.qry(dfn[u]));
    }
    return 0;
}