375-D 题目大意
给定一颗\(n\)个节点的树,根节点为\(1\),每个节点都有一个颜色\({c}_{i}\)
给定\(q\)次询问:
\((x,k):\)询问在以\(x\)为根的子树中出现次数\({\ge}k\)的颜色数量
Solution
首先考虑一般暴力的做法,将所有询问离线下来,对整个树进行\(dfs\),在\(dfs\)到节点\(x\)的时候,把相关\(x\)的询问统一处理了,具体处理如下。
预处理出\(dfs\)序,枚举\(x\)子树中的所有节点,用\(cnt\)数组统计颜色,\(sum\)数组统计颜色数量。递归返回时,清空\(cnt\)和\(sum\)两个数组。总复杂度为\(O(n^2+q)\)
显然,暴力的时间复杂度过高,考虑如何用\(dsu\,on\,tree\)优化。在枚举\(x\)的子树时是为了处理出\(cnt\)和\(sum\)两个数组。需要灵性的点在于观察到在\(dfs\)的过程中,枚举完\(x\)的最后一个儿子节点返回时可以不用清空这两个数组,而是使\(x\)继承这个儿子子树的信息,然后只需要去把\(x\)其他儿子子树的信息加上,就可以得到完成的\(x\)子树的颜色信息。
想到这里,接下来要考虑的是把哪个儿子节点作为继承人,显然,我们希望这个儿子子树的规模越大越好,这样能够重复利用的信息就更多。到这里整个算法的思路就显而易见了。
首先第一遍\(dfs\)预处理出每个节点的重儿子与整个树的\(dfn\)序,第二遍\(dfs\)优先遍历所有轻儿子,最后遍历重儿子,(先重后轻会清空两个记录颜色信息的数组,导致\(x\)节点继承不到任何信息)。遍历完之后,此时能够\(O(1)\)回答\(x\)所有的相关的询问,递归返回前,如果当前不是重儿子则清空两个数组。
这就是树上启发式合并的思想,可以证明时间复杂度为\(O(nlogn)\)
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N=1e5+10;
int col[N],sum[N];
void solve(){
int n,m;
cin>>n>>m;
vector<int> c(n+1);
for(int i=1;i<=n;i++) cin>>c[i];
vector<vector<int>> e(n+1);
for(int i=1;i<n;i++){
int x,y;
cin>>x>>y;
e[x].push_back(y);
e[y].push_back(x);
}
vector<vector<array<int,2>>> Q(n+1);
for(int i=0;i<m;i++){
int u,k;
cin>>u>>k;
Q[u].push_back({k,i});
}
vector<int> L(n+1),R(n+1),id(n+1);
int dfn=0;
vector<int> sz(n+1),son(n+1);
function<void(int,int)> dfs1=[&](int x,int fa){
L[x]=++dfn,sz[x]=1,id[dfn]=x;
for(auto y:e[x]){
if(y==fa) continue;
dfs1(y,x);
sz[x]+=sz[y];
if(sz[y]>sz[son[x]]) son[x]=y;
}
R[x]=dfn;
};
dfs1(1,0);
vector<int> ans(m);
function<void(int,int,int)> dfs2=[&](int x,int fa,int heavy){
for(auto y:e[x]){
if(y==fa||y==son[x]) continue;
dfs2(y,x,0);
}
if(son[x]) dfs2(son[x],x,1);
for(auto y:e[x]){
if(y==fa||y==son[x]) continue;
for(int i=L[y];i<=R[y];i++){
col[c[id[i]]]++,sum[col[c[id[i]]]]++;
}
}
col[c[x]]++,sum[col[c[x]]]++;
for(auto q:Q[x]){
ans[q[1]]=sum[q[0]];
}
if(!heavy){
for(int i=L[x];i<=R[x];i++){
sum[col[c[id[i]]]]--,col[c[id[i]]]--;
}
}
};
dfs2(1,0,1);
for(int i=0;i<m;i++) cout<<ans[i]<<" ";
}
int main(){
ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
int T=1;
//cin>>T;
while(T--){
solve();
}
return 0;
}