树上启发式合并/莫队 Tree and Queries

题目链接：Problem – D – Codeforces

题目大意：

一棵树上 (n <= 1e5) 每个点都有一个颜色 (ci <= 1e5) ，m (m <= 1e5) 次询问，每次询问包含一个 v 和一个 k ，统计 v 的子树中这样的颜色 x 的种类数：至少有 k 个点的颜色为 x 。

Solution1:

树上启发式合并，记录每个颜色出现的次数 col[i]，以及一个数组 num[i] 表示出现次数至少为 i 的颜色数。（注意我们可以在更新 col 的同时 O(1) 更新 num）

Code:

#include<cstdio>
#include<cstring>
#include<vector>
using namespace std;

#define N 100005

int n,m,cnt;
int c[N],st[N],siz[N],son[N],dfn[N],col[N],num[N],id[N],ans[N];

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

struct Edge
{
int next,to;
}ed[N << 1];

void add(int u,int v)
{
ed[++ cnt].next = st[u];
ed[cnt].to = v;
st[u] = cnt;
return;
}

void pdfs(int x,int fa)
{
siz[x] = 1;
dfn[x] = ++ cnt;
id[cnt] = x;
for (int i = st[x]; ~i ;i = ed[i].next)
{
int rec = ed[i].to;
if(rec == fa) continue;
pdfs(rec,x);
siz[x] += siz[rec];
if(!son[x] || siz[rec] > siz[son[x]]) son[x] = rec;
}
return;
}

void calc(int x) { ++ col[x] , ++ num[col[x]] ; }

void del(int x) { — num[col[x]] , — col[x] ; }

void dfs(int x,int fa,int op)
{
for (int i = st[x]; ~i ;i = ed[i].next)
{
int rec = ed[i].to;
if(rec == fa || rec == son[x]) continue;
dfs(rec,x,0);
}
if(son[x]) dfs(son[x],x,1);
calc(c[x]);
for (int i = st[x]; ~i ;i = ed[i].next)
{
int rec = ed[i].to;
if(rec == fa || rec == son[x]) continue;
for (int j = dfn[rec];j <= dfn[rec] + siz[rec] – 1;++ j)
calc(c[id[j]]);
}
for (auto i : q[x]) ans[i.first] = num[i.second];
if(!op)
{
for (int i = dfn[x];i <= dfn[x] + siz[x] – 1;++ i)
del(c[id[i]]);
}
return;
}

int main()
{
memset(st,-1,sizeof st);
scanf("%d%d",&n,&m),cnt = 0;
for (int i = 1;i <= n;++ i) scanf("%d",&c[i]);
for (int i = 1,u,v;i < n;++ i)
scanf("%d%d",&u,&v),add(u,v),add(v,u);
for (int i = 1,v,k;i <= m;++ i)
{
scanf("%d%d",&v,&k);
q[v].push_back({i,k});
}
cnt = 0;
pdfs(1,0);
dfs(1,0,1);
for (int i = 1;i <= m;++ i) printf("%d\\n",ans[i]);
return 0;
}

Solution2:

利用 "子树内节点的 dfn 是一段连续区间" ，可以将统计子树内的答案转化成统计区间内的答案，离线莫队处理。

Code:

#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;

#define N 100005

int n,m,cnt,B;
int c[N],st[N],dfn[N],siz[N],col[N],num[N],id[N],ans[N],bel[N];

struct Query
{
int l,r,bar,id;
}q[N];

struct Edge
{
int next,to;
}ed[N << 1];

int cmp(Query x,Query y)
{
if(bel[x.l] == bel[y.l]) return x.r < y.r ;
return bel[x.l] < bel[y.l] ;
}

void added(int u,int v)
{
ed[++ cnt].next = st[u];
ed[cnt].to = v;
st[u] = cnt;
return;
}

void dfs(int x,int fa)
{
dfn[x] = ++ cnt;
id[cnt] = x;
siz[x] = 1;
for (int i = st[x]; ~i ;i = ed[i].next)
{
int rec = ed[i].to;
if(rec == fa) continue;
dfs(rec,x);
siz[x] += siz[rec];
}
return;
}

void add(int x) { ++ col[c[id[x]]] , ++ num[col[c[id[x]]]] ; }

void sub(int x) { — num[col[c[id[x]]]] , — col[c[id[x]]] ; }

int main()
{
memset(st,-1,sizeof st);
scanf("%d%d",&n,&m),cnt = 0,B = (int)sqrt(n);
for (int i = 1;i <= n;++ i) scanf("%d",&c[i]),bel[i] = (i – 1) / B + 1;
for (int i = 1,u,v;i < n;++ i)
scanf("%d%d",&u,&v),added(u,v),added(v,u);
cnt = 0;
dfs(1,0);
for (int i = 1,v,k;i <= m;++ i)
{
scanf("%d%d",&v,&k);
q[i].l = dfn[v];
q[i].r = dfn[v] + siz[v] – 1;
q[i].bar = k;
q[i].id = i;
}
sort(q + 1,q + m + 1,cmp);
int l,r;
l = 1,r = 0;
for (int i = 1;i <= m;++ i)
{
while (l > q[i].l) add(– l);
while (l < q[i].l) sub(l ++);
while (r < q[i].r) add(++ r);
while (r > q[i].r) sub(r –);
ans[q[i].id] = num[q[i].bar];
}
for (int i = 1;i <= m;++ i) printf("%d\\n",ans[i]);
return 0;
}