[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;
}
}