Codeforces375D Tree and Queries

dsu on tree

题目链接

点我跳转

题目大意

给定一棵 \(n\) 个节点的树，根节点为 \(1\)。每个节点上有一个颜色 \(c_i\)

\(m\) 次询问。

每次询问给出 \(u\) \(k\)：询问在以 \(u\) 为根的子树中，出现次数 \(≥k\) 的颜色有多少种。

解题思路

可以开棵权值线段树

如果当前颜色出现的次数 \(cnt[i] = x\), 就把树的第 \(x\) 个位置的值 \(+ 1\)

那么对于每个询问的 \(k\) 输出树的第 \(k\) 个位置的值即可

AC_Code

#include<bits/stdc++.h>
#define rep(i,a,n) for (int i=a;i<=n;i++)
#define per(i,n,a) for (int i=n;i>=a;i--)
#define int long long
#define pb push_back
#define fi first
#define se second
using namespace std;
const int N = 3e5 + 10;
struct Tree{
	int l , r , lazy , sum;
}tree[N << 2];
void push_up(int rt)
{
	tree[rt].sum = tree[rt << 1].sum + tree[rt << 1 | 1].sum;
}
void push_down(int rt)
{
	int x = tree[rt].lazy;
	tree[rt].lazy = 0;
	tree[rt << 1].lazy = tree[rt << 1 | 1].lazy = x;
	tree[rt << 1].sum += (tree[rt << 1].r - tree[rt << 1].l + 1) * x;
	tree[rt << 1 | 1].sum += (tree[rt << 1 | 1].r - tree[rt << 1 | 1].l + 1) * x;
}
void build(int l , int r , int rt)
{
	tree[rt].l = l , tree[rt].r = r , tree[rt].lazy = 0;
	if(l == r)
	{
		tree[rt].sum = 0;
		return ;
	}
	int mid = l + r >> 1;
	build(l , mid , rt << 1);
	build(mid + 1 , r , rt << 1 | 1);
	push_up(rt);
}
void update_range(int L , int R , int rt , int val)
{
	int l = tree[rt].l , r = tree[rt].r;
	if(L <= l && r <= R)
	{
		tree[rt].lazy += val;
		tree[rt].sum += (r - l + 1) * val;
		return ;
	}
	push_down(rt);
	int mid = l + r >> 1;
	if(L <= mid) update_range(L , R , rt << 1 , val);
	if(R > mid)  update_range(L , R , rt << 1 | 1 , val);
	push_up(rt);
}
int query_range(int L , int R , int rt)
{
	int l = tree[rt].l , r = tree[rt].r;
	if(L <= l && r <= R) return tree[rt].sum;
	push_down(rt);
	int mid = l + r >> 1 , ans = 0;
	if(L <= mid) ans += query_range(L , R , rt << 1);
	if(R > mid)  ans += query_range(L , R , rt << 1 | 1);
	return ans;
}
struct Edge{
	int nex , to;
}edge[N << 1];
int head[N] , TOT;
void add_edge(int u , int v)
{
	edge[++ TOT].nex = head[u] ;
	edge[TOT].to = v;
	head[u] = TOT;
}
int dep[N] , sz[N] , hson[N] , HH;
int col[N] , n , m , up;
int cnt[N] , sum[N];
vector<pair<int , int>>Q[N] , ans;
void dfs(int u , int far)
{
	dep[u] = dep[far] + 1;
	sz[u] = 1;
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far) continue ;
		dfs(v , u);
		sz[u] += sz[v];
		if(sz[v] > sz[hson[u]]) hson[u] = v;
	}
}
void calc(int u , int far,  int val)
{
	cnt[col[u]] += val;
	if(val == 1)
	{
		int k = cnt[col[u]];
		update_range(k , k , 1 , 1);
	}
	if(val == -1)
	{
		int k = cnt[col[u]] + 1;
		update_range(k , k , 1 , -1);
	}
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == HH) continue ;
		calc(v , u , val);
	}
}
void dsu(int u , int far , int op)
{
	for(int i = head[u] ; i ; i = edge[i].nex)
	{
		int v = edge[i].to;
		if(v == far || v == hson[u]) continue ;
		dsu(v , u , 0);
	}
	if(hson[u]) dsu(hson[u] , u , 1) , HH = hson[u];
	calc(u , far , 1);
	for(auto i : Q[u])
	{
		int id = i.fi , k = i.se;
		int res = query_range(k , k , 1);
		ans.pb(make_pair(id , res));
	}
	HH = 0;
	if(!op) calc(u , far , -1);
}
signed main()
{
	ios::sync_with_stdio(false);
	cin.tie(0) , cout.tie(0);
	cin >> n >> m;
	rep(i , 1 , n) cin >> col[i];
	rep(i , 1 , n - 1)
	{
		int u , v;
		cin >> u >> v;
		add_edge(u , v) , add_edge(v , u);
	}
	rep(i , 1 , m)
	{
		int u , k;
		cin >> u >> k;
		Q[u].pb(make_pair(i , k));
	}
	build(1 , 100000 , 1);
	dfs(1 , 0);
	dsu(1 , 0 , 0);
	sort(ans.begin() , ans.end());
	for(auto i : ans) cout << i.se << '\n';
	return 0;
}