P8339 [AHOI2022] 钥匙

题意

给定一颗树，树上每个节点都有一把钥匙 $key$ 或者一个宝箱 $box$，对应颜色的钥匙才能打开对应的宝箱，钥匙只能使用一次，相同颜色的钥匙最多只有 $5$ 把。

现在有若干次询问，每次询问给定一组 $(s,e)$，求从 $s$ 走到 $e$ 能打开多少个宝箱。

题解

精简一下题目，这个问题可以理解为树上的括号匹配问题，下面称 $(key,box)$ 为一组匹配，即用这个 $key$ 可以沿着某个路径打开这个 $box$。

乍一看很难做，这个时候就得揪着很特殊的条件分析。

没错，特殊条件就是相同颜色的钥匙最多只有 $5$ 把，所以匹配数量是 $O(n)$ 的，可以将所有的匹配预处理出来。

实现方法：将相同颜色地节点提取出来单独处理，使用了虚树，分离出来后由于同种颜色地钥匙最多只有 $5$ 把，所以我们可以分别以这几把钥匙为根进行 $\text{dfs}$ 搜索，找到匹配的条件可以直接理解为括号匹配，想必这样子就不必多说了。

这样一来每次询问就等价于统计一段路径上的匹配数量。

首先如果路径 $(key,box)$ 只有被路径 $(s,e)$ 覆盖了，才会对其产生贡献。

由于很难直接针对每个询问分别统计，于是考虑一个匹配会对哪些询问 $(s,e)$ 产生贡献，分三种情况：

1. $box$ 是 $key$ 的祖先： $s$ 要在 $key$ 的子树中，$e$ 不能在路径 $(key,box)$ 上与 $box$ 相邻的节点的子树中。

2. $key$ 是 $box$ 的祖先： 跟 1. 同理。

3. $key$ 和 $box$ 没有祖先关系： $s$ 要在 $key$ 的子树中，$e$ 要在 $box$ 的子树中。

很自然的想到区间的概念，我们需要把这抽象的条件转换为区间，从子树得到灵感，无非就是转化为 $\text{dfn}$ 序，因为 $\text{dfn}$ 序在子树内是连续的。

所以我们得到了这个信息，对于一对匹配 $(key,box)$，想要得到其贡献，$s$ 和 $e$ 的 $\text{dfn}$ 序必须在某个区间内。

我们以 $s$ 的 $\text{dfn}$ 序为横轴，$e$ 的 $\text{dfn}$ 序为竖轴，那么问题就转化为用矩形覆盖坐标系，再单点查询被覆盖了几次，用区间修改，单点查询的方法解决。（代码第 $223$ 行到第 $234$ 行）

实现方法：由于是平面上，所以我们使用差分+树状数组+二位偏序实现，具体地，将询问离线，再用下标的形式降掉一维，再用权值树状数组维护另外一维统计答案。

最后，由于要算 $\text{dfn}$ 序和 $\text{LCA}$（虚树中要用），所以可以直接打一个重链剖分。

总结：由特殊条件想到预处理可能产生的贡献 $(key,box)$，再根据每对 $(key,box)$ 的作用范围进行区间覆盖，单点查询。全程保证了总时间复杂度为 $O((n+m) \log n)$。

代码

#include <bits/stdc++.h>
#define fi first
#define se second
#define eb emplace_back

using namespace std;
using pii = pair<int,int>;

const int N=5e5+5;
const int M=1e6+5;

int n,m;
int t[N],c[N];
vector<int> e[N],E[N],col[N];
vector<pii> q[N];

//树链剖分
int fa[N],dep[N],siz[N],son[N];
void dfs1(int u){
	siz[u]=1;
	for(auto v:e[u]){
		if(v==fa[u]) continue;
		fa[v]=u;
		dep[v]=dep[u]+1;
		dfs1(v);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]]) son[u]=v;
	}
}

int dfn[N],top[N],tot;
void dfs2(int u,int t){
	top[u]=t;
	dfn[u]=++tot;
	if(!son[u]) return;
	dfs2(son[u],t);
	for(auto v:e[u]) if(v!=fa[u]&&v!=son[u])  dfs2(v,v);
}

void slpf() {dfs1(1),dfs2(1,1);}
/*------*/

//LCA
int LCA(int x,int y){
	while(top[x]!=top[y]){
		if(dep[top[x]]<dep[top[y]]) swap(x,y);		
		x=fa[top[x]]; 
	}
	return dep[x]<dep[y]?x:y;
}

//求路径上的次端点（对的可以这样理解）
int find(int x,int y){
	if(dep[x]<dep[y]) swap(x,y);
	int lst;
	while(top[x]!=top[y]) lst=top[x],x=fa[top[x]];
	if(x==y) return lst;
	return son[y];
}

//平面差分
// dfn[s] in [a, b], dfn[e] in [c, d]: +1
vector<pii> row[N];
void add(int a,int b,int c,int d){
	if(a>b||c>d) return;
	row[a].eb    (c,1);
	row[a].eb  (d+1,-1);
	row[b+1].eb  (c,-1);
	row[b+1].eb(d+1,1);
}

//将 (key,box) 插入dfn平面
void insert(int x,int y){
	int l=LCA(x,y);
	if(l==x){
		int z=find(x,y);
		//警示后人：这种很多逗号的地方一定要对齐，要不然你调半天都看不出来这里打错了。
		add(1            ,dfn[z]-1,dfn[y],dfn[y]+siz[y]-1);
		add(dfn[z]+siz[z],n       ,dfn[y],dfn[y]+siz[y]-1);
	}else if(l==y){
		int z=find(x,y);
		add(dfn[x],dfn[x]+siz[x]-1,1            ,dfn[z]-1);
		add(dfn[x],dfn[x]+siz[x]-1,dfn[z]+siz[z],n);
	}else add(dfn[x],dfn[x]+siz[x]-1,dfn[y],dfn[y]+siz[y]-1);
}


//有关 找出 (key,box) 配对的 函数。
namespace keybox{
	//Virtual_Tree：虚树
	vector<int> id;
	int a[N],m,b[N<<1],len;
	bool cmp (int x,int y) {return dfn[x]<dfn[y];}
	void clear(){
		m=len=0;
		while(!id.empty()){
			E[id.back()].clear();
			id.pop_back();
		}
	}
	//建虚树（二次排序 + LCA 连边）
	void build_VT(int c){
		clear();
		for(int x:col[c]) a[++m]=x;
		sort(a+1,a+m+1,cmp);
		for(int i=1;i<m;i++){
			b[++len]=a[i];
			b[++len]=LCA(a[i],a[i+1]);
		}
		b[++len]=a[m];
		sort(b+1,b+len+1,cmp);
		len=unique(b+1,b+len+1)-b-1;
		for(int i=1;i<len;i++){
			int l=LCA(b[i],b[i+1]);
			id.eb(l),id.eb(b[i+1]);
			E[l].eb(b[i+1]),E[b[i+1]].eb(l);
		}
	}
	//以key为起点找与其配对的box，dfs最多被调用5次
	void dfs(int u,int fa,int w,int col,int rt){
		if(c[u]==col){
			if(t[u]==1) w++;
			else{
				w--;
				if(!w){
					insert(rt,u);
					return;
				}
			}
		}
		for(auto v:E[u]) if(v!=fa) dfs(v,u,w,col,rt);
	}
	void search(int col){
        for(int i=1;i<=m;i++) 
            if(c[a[i]]==col&&t[a[i]]==1) 
                dfs(a[i],a[i],0,col,a[i]);
    }
}

//树状数组
namespace BIT{
	int t[N];
	void upd(int x,int y){
		for(;x<N;x+=x&-x) t[x]+=y;
	}
	int qry(int x){
		int res=0;
		for(;x;x-=x&-x) res+=t[x];
		return res;
	}
}

int ans[M];

int main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

	cin>>n>>m;
	for(int i=1;i<=n;i++) cin>>t[i]>>c[i],col[c[i]].eb(i);
	for(int i=1;i<n;i++){
		int u,v;
		cin>>u>>v;
		e[u].eb(v),e[v].eb(u);
	}

	slpf();

	for(int i=1;i<=n;i++){
		if(col[i].empty()) continue;
		keybox::build_VT(i);
		keybox::search(i);
	}

    //将询问离线
	for(int i=1;i<=m;i++){
		int s,e;
		cin>>s>>e;
		q[dfn[s]].eb(dfn[e],i);
	}

    //树状数组+二维偏序
	for(int i=1;i<=n;i++){
		for(auto it:row[i]) BIT::upd(it.fi,it.se);
		for(auto it:q[i]) ans[it.se]=BIT::qry(it.fi);
	}

	for(int i=1;i<=m;i++) cout<<ans[i]<<"\n";
    return 0;
}