CodeForces - 375D Tree and Queries 树上启发式合并 树状数组

CodeForces - 375D Tree and Queries

题意

一棵 $n$ 个节点的树，每个节点都有一个颜色 $co{l}_i$ ，有 $m$ 次询问，问以 $v$ 为根的子树中，有多少个颜色 $c$ 满足至少有 $k$ 个点的颜色为 $c$ 。

解法

离线的子树问题，考虑树上启发式合并（dsu on tree​）。

• 每次统计出子树中某个颜色的数量，然后用树状数组维护颜色数量为某个值的答案。
• 使用 dsu on tree 来优化搜索，使得每次先搜轻孩子，删去贡献后搜重孩子，并保留贡献。

解法

#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e5 + 5;
int n, q;
int a[maxn];
vector<int> g[maxn];
struct node {
    int id, k;
    node() {}
    node(int id, int k) : id(id), k(k) {}
};
vector<node> Q[maxn];
int sz[maxn], son[maxn];
int cnt[maxn], res[maxn], flag;
int cntt[maxn];
void dfs1(int u, int f) {
    sz[u] = 1;
    for (auto v : g[u]) {
        if (v == f) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        if (sz[v] > sz[son[u]]) son[u] = v;
    }
}
int c[maxn];
void add(int x, int val) {
    if (!x) return;
    while (x < maxn) c[x] += val, x += x & -x;
}
int getsum(int x) {
    int ans = 0;
    while (x) ans += c[x], x -= x & -x;
    return ans;
}
void count(int u, int f, int val) {
    add(cnt[a[u]], -1);
    cnt[a[u]] += val;
    add(cnt[a[u]], 1);
    for (auto v : g[u]) {
        if (v == f || v == flag) continue;
        count(v, u, val);
    }
}
void query(int u) {
    for (auto e : Q[u]) {
        res[e.id] = getsum(maxn - 1) - getsum(e.k - 1);
    }
}
void dfs(int u, int f, bool keep) {
    for (auto v : g[u]) {
        if (v == f || v == son[u]) continue;
        dfs(v, u, 0);
    }
    if (son[u]) {
        dfs(son[u], u, 1);
        flag = son[u];
    }
    count(u, f, 1);
    query(u);
    flag = 0;
    if (!keep) count(u, f, -1);
}
int main() {
    scanf("%d%d", &n, &q);
    rep(i, 1, n) scanf("%d", &a[i]);
    rep(i, 1, n - 1) {
        int u, v;
        scanf("%d%d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    rep(i, 1, q) {
        int u, k;
        scanf("%d%d", &u, &k);
        Q[u].push_back(node(i, k));
    }
    dfs1(1, 0);
    dfs(1, 0, 0);
    rep(i, 1, q) printf("%d\n", res[i]);
}