[NOISG2022 Qualification] Dragonfly Solution in O(d log d)

[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;
}