[CF375D] Tree and Queries

[CF375D] Tree and Queries - dsu on tree

Description

给定一棵树，根为 (1)，每个点上有颜色 (c_i)，有 (m) 次询问，每次指定 (u,k)，问 (u) 为根的子树中，出现次数 (ge k) 的颜色有多少种。

Solution

dsu on tree，每次找到先处理所有轻儿子，最后处理重儿子，处理完重儿子后不清空，暴力加上所有轻儿子的贡献，得到该子树内所有点的贡献。

由于要维护的是有多少颜色数 (ge k)，我们很容易想到用 BIT 的做法，但其实并不需要。

设 (cnt[i]) 表示颜色 (i) 出现了多少次，设 (sum[i]) 表示出现次数大于等于 (i) 的颜色数量有多少，那么在每次假设现在 (cnt[i]) 从 (x) 变为 (x+1)，我们只需要在 (sum[x+1]) 上加一即可，反之亦然。

#include <bits/stdc++.h>
using namespace std;

#define int long long
const int N = 1000005;

int n, m;
vector<int> g[N];
int a[N], siz[N], wson[N], ans[N];

// bucket

int cnt[N], sum[N];
vector<int> buf_opera;

vector<pair<int, int>> ask[N];

void add(int p)
{
    buf_opera.push_back(p);
    cnt[p]++;
    sum[cnt[p]]++;
}

void dec(int p)
{
    sum[cnt[p]]--;
    cnt[p]--;
}

void clear()
{
    while (buf_opera.size())
    {
        int p = buf_opera.back();
        dec(p);
        buf_opera.pop_back();
    }
}

// calc size and wson

void calc(int p, int from)
{
    siz[p] = 1;
    for (int q : g[p])
    {
        if (q != from)
        {
            calc(q, p);
            siz[p] += siz[q];
            if (siz[q] > siz[wson[p]])
                wson[p] = q;
        }
    }
}

void add_r(int p, int from)
{
    add(a[p]);
    for (int q : g[p])
    {
        if (q != from)
        {
            add_r(q, p);
        }
    }
}

// dfs solve

void dfs(int p, int from)
{
    for (int q : g[p])
    {
        if (q != from && q != wson[p])
        {
            dfs(q, p);
            clear();
        }
    }
    if (wson[p])
    {
        dfs(wson[p], p);
    }
    add(a[p]);
    for (int q : g[p])
    {
        if (q != from && q != wson[p])
        {
            add_r(q, p);
        }
    }
    for (auto [id, k] : ask[p])
    {
        ans[id] = sum[k];
    }
    // for (int i = 1; i <= n; i++)
    // {
    //     cout << sum[i] << ",";
    // }
    // cout << endl;
}

// main

signed main()
{
    ios::sync_with_stdio(false);

    cin >> n >> m;
    for (int i = 1; i <= n; i++)
        cin >> a[i];
    for (int i = 1; i < n; i++)
    {
        int t1, t2;
        cin >> t1 >> t2;
        g[t1].push_back(t2);
        g[t2].push_back(t1);
    }

    for (int i = 1; i <= m; i++)
    {
        int u, k;
        cin >> u >> k;

        ask[u].push_back({i, k});
    }

    calc(1, 0);
    dfs(1, 0);

    for (int i = 1; i <= m; i++)
    {
        cout << ans[i] << endl;
    }
}