P11828 [TOIP2024] 距离函数 题解
Qing_Nian
·
·
题解
简要题意
给你两棵 n 个点的有根树，定义 dis(u,v)=\min(dep_{u}-dep_{\mathrm{lca}(u,v)},dep_{v}-dep_{\mathrm{lca}(u,v)})。
计算在满足一棵树中 dis(u,v)=0 在另一棵树中 dis(u,v)\gt k 的无序点对数量。
$1\le n\le 2\times 10^5$。
### 思路
下面记 $dfn_u$ 为点 $u$ 的dfn 序，$size_u$ 为点 $u$ 的子树大小，称 $(u,v)$ 存在祖先关系为其中一个为另一个在树中的祖先。
首先 $dis(u,v)=0$ 等价于点对 $(u,v)$ 存在祖先关系。
考虑如何快速判断 $(u,v)$ 是否存在祖先关系。
假设 $dep_u\lt dep_v$。
则 $(u,v)$ 有祖先关系的充要条件为 $[dfn_v,dfn_v+size_v-1]\subseteq [dfn_u,dfn_u+size_u-1]$，下面简记 $dfn_u$ 为 $L_u$，$dfn_u+size_u-1$ 为 $R_u$。
- $k=0$ 时
$(u,v)$ 在一棵树上存在祖先关系，在另一棵树上不存在祖先关系。
$(u,v)$ 不存在祖先关系即 $[L_u,R_u]\cap [L_v,R_v]= \emptyset$。
拆一下即为找满足 $L_u\gt R_v$ 或 $L_v\gt R_u$ 的点对，二维数点模板，我们开两个树状数组。
在遍历一棵树时，将另一棵树的 $L,R$ 值分别丢进树状数组维护即可，具体维护方式可以参考代码。
- $k\gt 0$ 时
$(u,v)$ 的距离 $\gt k$ 等价于的双方的 $k$ 级祖先不存在祖先关系，那么我们按 $k=0$ 时的方式维护，维护的每个点的 $L,R$ 值变为每个点 $k$ 级祖先的 $L,R$ 值即可。
### Code
```cpp
/*
*/
#include <bits/stdc++.h>
using namespace std;
inline long long read(){
long long x=0,w=0;char c=0;
while(!isdigit(c)) {w|=c=='-';c=getchar();}
while(isdigit(c)) x=(x<<3)+(x<<1)+(c^48),c=getchar();
return w?-x:x;
}
int n,k;
int root[2];
long long ans;
vector<int> mp[2][200050];
int f[2][20][200050]；
int dep[2][200050],L[2][200050],R[2][200050],num;//与题解中定义相同
void dfs(int ty,int x,int fa){
f[ty][0][x]=fa;
for(int i=1;i<20;i++){
f[ty][i][x]=f[ty][i-1][f[ty][i-1][x]];
}
dep[ty][x]=dep[ty][fa]+1;
L[ty][x]=++num;
for(int v:mp[ty][x]){
dfs(ty,v,x);
}
R[ty][x]=num;
}
struct BIT{
long long c[200050];
inline int lowbit(int x){
return x&(-x);
}
inline void up(int x,int y){
while(x<=n){
c[x]+=y;
x+=lowbit(x);
}
}
inline int query(int x){
int res=0;
while(x){
res+=c[x];
x-=lowbit(x);
}
return res;
}
}TL,TR;
int kx[2][200050];//存每个点的 k 级祖先
void dfs2(int ty,int x){
if(kx[ty^1][x]){//排除 0 防止树状数组暴毙
ans=ans+TL.query(-R[ty^1][kx[ty^1][x]]+n);
ans=ans+TR.query(L[ty^1][kx[ty^1][x]]-1);
TR.up(R[ty^1][kx[ty^1][x]],1);
TL.up(-L[ty^1][kx[ty^1][x]]+n+1,1);
}
for(int v:mp[ty][x]){
dfs2(ty,v);
}
if(kx[ty^1][x]){//递归结束撤销，在点 v 时树状数组中有它所有祖先的信息
TR.up(R[ty^1][kx[ty^1][x]],-1);
TL.up(-L[ty^1][kx[ty^1][x]]+n+1,-1);
}
}
int main(){
n=read();
k=read();
for(int i=1,x;i<=n;i++){
x=read();
if(!x){
root[0]=i;
}
else{
mp[0][x].push_back(i);
}
}
for(int i=1,x;i<=n;i++){
x=read();
if(!x){
root[1]=i;
}
else{
mp[1][x].push_back(i);
}
}
dfs(0,root[0],0);
num=0;
dfs(1,root[1],0);
for(int ty=0,x;ty<2;ty++){
for(int i=1;i<=n;i++){
x=k;
kx[ty][i]=i;
for(int j=19;~j;j--){
if(x>=(1<<j)){
x-=(1<<j);
kx[ty][i]=f[ty][j][kx[ty][i]];//倍增求 k 级祖先
}
}
}
}
dfs2(0,root[0]);
dfs2(1,root[1]);
printf("%lld",ans);
return 0;
}
/*
*/
```