CodeForces 375D. Tree and Queries【树上启发式合并】
题意
给出一棵 \(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\n",ans[i]);
return 0;
}