CodeForces 375D. Tree and Queries树上启发式合并

Posted 2021-03-02 bakacirno

传送门

题意

给出一棵 (n) 个结点的树，每个结点有一个颜色 (c_i) 。 询问 (q) 次，每次询问以 (v) 结点为根的子树中，出现次数 (ge k) 的颜色有多少种。树的根节点是 (1)。

题解

反正我看见这个 (ge k) 就觉得要用线段树，实际上好像不用写线段树的 Orz。
还是树上启发式合并，记录每种颜色出现的次数，然后线段树记录某种次数有多少颜色，更改就在线段树上改。
这是最后一道树上启发式合并的例题了，以后遇到再刷。

#include <bits/stdc++.h>
#define xx first
#define yy second
using namespace std;
typedef pair<int,int> PII;
const int N=1e5+10;
int n,m,clr[N],siz[N],son[N],cnt[N],ans[N],Son;
vector<PII> ask[N];
vector<int> g[N];
struct SegTree{
	#define mid (l+r>>1)
	int sum[N*4];
	void upd(int id,int l,int r,int pos,int x){
		if(l==r) {sum[id]+=x;return;}
		if(pos<=mid) upd(id<<1,l,mid,pos,x);
		else upd(id<<1|1,mid+1,r,pos,x);
		sum[id]=sum[id<<1]+sum[id<<1|1];
	}
	int ask(int id,int l,int r,int L,int R){
		if(L<=l&&r<=R) return sum[id];
		int res=0;
		if(L<=mid) res+=ask(id<<1,l,mid,L,R);
		if(R>mid) res+=ask(id<<1|1,mid+1,r,L,R);
		return res;
	}
	#undef mid
}tr;

void predfs(int u,int fa){
	siz[u]=1;
	for(int v:g[u]){
		if(v==fa) continue;
		predfs(v,u);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}

void add(int u,int fa,int val){
	tr.upd(1,0,1e5,cnt[clr[u]],-1);
	cnt[clr[u]]+=val;
	tr.upd(1,0,1e5,cnt[clr[u]],1);
	for(int v:g[u]) if(v!=fa&&v!=Son) add(v,u,val);
}

void dfs(int u,int fa,bool keep){
	for(int v:g[u]){
		if(v==fa||v==son[u]) continue;
		dfs(v,u,false);
	}
	if(son[u]) dfs(son[u],u,true),Son=son[u];
	add(u,fa,1);
	for(PII q:ask[u]) ans[q.xx]=tr.ask(1,0,1e5,q.yy,1e5);
	if(!keep) Son=0,add(u,fa,-1);
}

int main(){
	scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++) scanf("%d",&clr[i]);
	for(int i=1,u,v;i<n;i++){
		scanf("%d%d",&u,&v);
		g[u].push_back(v);
		g[v].push_back(u);
	}
	for(int i=1,u,k;i<=m;i++){
		scanf("%d%d",&u,&k);
		ask[u].push_back({i,k});
	}
	predfs(1,0);
	tr.upd(1,0,1e5,0,1e5);
	dfs(1,0,true);
	for(int i=1;i<=m;i++) printf("%d
",ans[i]);
	return 0;
}