[NOISG2022 Qualification] Dragonfly Solution in O(d log d)
提供一个使用线段树合并、栈、树状数组的严格单 \(\log\) 离线做法。
题目大意：给你一棵树，每个点有权值和颜色，每次问你一个从 \(1\) 开始的路径，求权值不为 \(0\) 的节点的颜色种类数，并且把所有权值不为 \(0\) 的节点权值减一，询问之间不独立。
解题思路：考虑离线下来，因为一个节点权值为 \(0\) 后就没有贡献了，那么考虑求出最后能产生贡献的时刻 \(t_i\)。
那么能吃节点 \(i\) 的询问在 \(i\) 的子树内。考虑以时间为下标建线段树，询问就在这个询问所在的时间插入一个点，那么我们可以自底向上线段树合并，合并到点 \(i\) 时在线段树上二分即可求出 \(t_i\)。这一部分的线段树只需维护一个子树大小。这一部分是 \(O(d\log d)\) 的。
求出 \(t_i\) 后，对于第 \(j\) 个询问，我们就是求满足下列条件的颜色个数：从 \(1\) 到 \(h_j\) 路径上存在这种颜色的点 \(i\) 使得 \(t_i\ge j\)。
考虑从根开始遍历，进入一个节点就插入这个节点的贡献，离开子树就撤销贡献。
我们现在就想知道从 \(1\) 到当前点的路径上，每种颜色 \(t\) 的最大值是多少。然后每次最大值更新时用数据结构维护一下即可。
由于如果存在颜色相同的 \(i,j\) 满足 \(i\) 是 \(j\) 的祖先，且 \(t_i\ge t_j\)，那么 \(j\) 是没有用的。所以考虑对每种颜色开一个栈，当 \(t\) 大于栈顶时才推进栈中，然后用一个以 \(t\) 为下标的树状数组维护：有多少种颜色的 \(t\) 大于等于某个数。
这一部分也是 \(O(d\log d)\) 的。
然后就做完了，最终的复杂度就是 \(O(d\log d)\) 的（这里我们设 \(n,d\) 同阶）。
实现细节：
注意到线段树合并时，我们动态开点的静态数组可能会 MLE。但是线段树合并时很多合并后的节点就扔掉了。考虑用一个 vector 回收节点即可。
另外注意 \(n,d\) 的规模是不一样的，数组不要开小了。
AC 代码 ：
const int N=2e5+5,M=2e6+5,L=2e7;
int tot,ls[L],rs[L],sz[L],rt[N];
vi emp;
#define mid ((l+r)>>1)
void merge(int &x,int y,int l,int r) {
if(!x) {x=y; return;}
if(!y) {return;}
if(l==r) {sz[x]+=sz[y]; return;}
merge(ls[x],ls[y],l,mid),merge(rs[x],rs[y],mid+1,r);
sz[x]=sz[ls[x]]+sz[rs[x]];
emp.pb(y);
}
void add(int &x,int l,int r,int y) {
if(!x) {
if(tot<L-1) x=++tot;
else x=Back(emp),emp.pop_back(),sz[x]=ls[x]=rs[x]=0;
}
if(l==r) {sz[x]++; return;}
if(y<=mid) add(ls[x],l,mid,y);
else add(rs[x],mid+1,r,y);
sz[x]=sz[ls[x]]+sz[rs[x]];
}
int query(int x,int l,int r,int y) {
if(l==r) return l;
if(sz[ls[x]]>=y) return query(ls[x],l,mid,y);
return query(rs[x],mid+1,r,y-sz[ls[x]]);
}
struct BIT{
#define lowbit(x) (x&(-x))
int s[M];
void update(int x,int y) {
while(x<M) s[x]+=y,x+=lowbit(x);
}
int query(int x){
int _s=0;
while(x) _s+=s[x],x-=lowbit(x);
return _s;
}
}bit;
vi g[N],q[N];
int n,Q;
int b[N],s[N],h[M],t[N];
void dfs1(int x,int y) {
for(int v:g[x]) {
if(v==y) continue;
dfs1(v,x);
merge(rt[x],rt[v],1,Q);
}
for(int v:q[x]) add(rt[x],1,Q,v);
if(b[x]) {
if(sz[rt[x]]<b[x]) t[x]=Q;
else t[x]=query(rt[x],1,Q,b[x]);
}
}
stack<int> mx[N];
int ans[M];
void solve(int x,int y){
int co=s[x],flag=0;
if(t[x]) {
if(mx[co].empty()||t[x]>mx[co].top()) {
if(Size(mx[co])) bit.update(M-mx[co].top(),-1);
mx[co].push(t[x]);
bit.update(M-t[x],1);
flag=1;
}
}
for(int v:q[x]) {
ans[v]=bit.query(M-v);
}
for(int v:g[x]) {
if(v==y) continue;
solve(v,x);
}
if(flag) {
bit.update(M-t[x],-1);
mx[co].pop();
if(Size(mx[co])) bit.update(M-mx[co].top(),1);
}
}
signed main(){
read(n,Q);
fo(i,1,n) read(b[i]);
fo(i,1,n) read(s[i]);
fo(i,1,Q) read(h[i]),q[h[i]].pb(i);
fu(i,1,n) {
int u,v; read(u,v);
g[u].pb(v),g[v].pb(u);
}
dfs1(1,0);
solve(1,0);
fo(i,1,Q) write(ans[i],' ');
return 0;
}