Codeforces 375D dsu on tree + BIT

题意

传送门 Codeforces 375D Tree and Queries

题解

$dsuontree$ 统计子树信息。以颜色的数量为值域，使用 $BIT$ 维护颜色的个数。为了方便处理 $cnt[color]=0$ 或 $k>n$ 这样的边界情况，将值域翻转，即 $i$ 在 $BIT$ 中的索引为 $maxn-1-i$。总时间复杂度 $O(N{\log }^{2}N)$。

#include <bits/stdc++.h>
using namespace std;
typedef pair<int, int> P;
const int maxn = 100005, maxm = 100005, maxe = 2 * maxn;
int N, M, C[maxn], res[maxm], sz[maxn], cnt[maxn];
vector<P> Q[maxn];
int tot, head[maxn], to[maxe], nxt[maxe];
bool hv[maxn];
int bit[maxn];

int sum(int i)
{
    i = maxn - 1 - i;
    int s = 0;
    while (i)
        s += bit[i], i -= i & -i;
    return s;
}

void add(int i, int x)
{
    i = maxn - 1 - i;
    while (i < maxn)
        bit[i] += x, i += i & -i;
}

inline void add_edge(int x, int y) { to[++tot] = y, nxt[tot] = head[x], head[x] = tot; }

void get_sz(int x, int f)
{
    sz[x] = 1;
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f)
            get_sz(y, x), sz[x] += sz[y];
    }
}

void add_t(int x, int f, int d)
{
    int &c = cnt[C[x]];
    add(c, -1), add(c += d, 1);
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f && !hv[y])
            add_t(y, x, d);
    }
}

void dfs(int x, int f, int keep)
{
    int hs = 0;
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f && sz[y] > sz[hs])
            hs = y;
    }
    for (int i = head[x]; i; i = nxt[i])
    {
        int y = to[i];
        if (y != f && y != hs)
            dfs(y, x, 0);
    }
    if (hs)
        dfs(hs, x, 1), hv[hs] = 1;
    add_t(x, f, 1);
    if (hs)
        hv[hs] = 0;
    for (P &p : Q[x])
        res[p.second] = sum(p.first);
    if (!keep)
        add_t(x, f, -1);
}

int main()
{
    scanf("%d%d", &N, &M);
    for (int i = 1; i <= N; ++i)
        scanf("%d", C + i);
    for (int i = 1, x, y; i < N; ++i)
        scanf("%d%d", &x, &y), add_edge(x, y), add_edge(y, x);
    for (int i = 1, x, k; i <= M; ++i)
        scanf("%d%d", &x, &k), Q[x].push_back(P(k, i));
    get_sz(1, 0);
    dfs(1, 0, 1);
    for (int i = 1; i <= M; ++i)
        printf("%d\n", res[i]);
    return 0;
}