树上启发式合并习题

树上启发式合并（dsu on tree）

算法流程

首先遍历轻儿子，遍历完之后将数据清空。然后遍历所有重儿子数据不清空
此时继续遍历轻儿子，此时就可以统计答案了
最后根据：当前 u 是否为 fa 的轻儿子，来决定是否保留数据。如果是轻儿子就清空
清空的时候，都是先清空轻儿子的信息，然后会保留重儿子的信息，后面又会重新遍历轻儿子，因而可以保证数据正确。
时间复杂度：一个点距离根节点，最多经过 $log_2n$ 条轻边，因此最多被访问 $1+log_2n$ 次，加 1 是因为每个点首先会被访问一次。总时间复杂度： $nlog_2n$

习题

CF600E. Lomsat gelral

链接：https://codeforces.com/contest/600/problem/E

题意：给定一棵以 1 为根的树，每个节点有颜色。问以每个节点为根的子树中，颜色出现次数最多的编号之和。（次数最多的可能有多个颜色）

思路：更新的时候，记录颜色出现的次数，次数超过或者等于最大值时，做相应的更新操作

写法一

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=1e5+5;

int n;
vector<int> e[maxn];
int col[maxn],sz[maxn],son[maxn];
ll ans[maxn],cnt[maxn],sum,mx;

void update(int u,int fa,int val,int heavy)
{
    cnt[col[u]]+=val;
    if(cnt[col[u]]>mx) sum=col[u],mx=cnt[col[u]];
    else if(cnt[col[u]]==mx) sum+=col[u];
    for(auto v: e[u])
    {
        if(v==fa||v==heavy) continue;
        update(v,u,val,heavy);
    }
}
void dfs1(int u,int fa)
{
    sz[u]=1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
}
void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    update(u,fa,1,son[u]);
    ans[u]=sum;
    if(!kp) update(u,fa,-1,-1),sum=0,mx=0;
}
int main()
{
    scanf("%d",&n);
    for(int i=1; i<=n; ++i) scanf("%d",&col[i]);
    for(int i=1; i<=n-1; ++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    dfs1(1,0);
    dfs2(1,0,1);
    for(int i=1; i<=n; ++i) printf("%lld%c",ans[i],i==n?'\n':' ');
    return 0;
}

写法二

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=1e5+5;
int n,col[maxn];
int in[maxn],out[maxn],id[maxn],times;
int sz[maxn],son[maxn];
ll mx,sum,cnt[maxn],ans[maxn];
vector<int> e[maxn];
void add(int u,int f)
{
    cnt[col[u]]+=f;
    if(cnt[col[u]]>mx) mx=cnt[col[u]],sum=col[u];
    else if(cnt[col[u]]==mx) sum+=col[u];
}
void dfs1(int u,int fa)
{
    in[u]=++times;
    id[times]=u;
    sz[u]=1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
    out[u]=times;
}
void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    add(u,1);
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        for(int i=in[v]; i<=out[v]; ++i) add(id[i],1);
    }
    ans[u]=sum;
    if(!kp) for(int i=in[u]; i<=out[u]; ++i) add(id[i],-1),mx=0,sum=0;
}
int main()
{
    scanf("%d",&n);
    for(int i=1; i<=n; ++i) scanf("%d",&col[i]);
    for(int i=1; i<=n-1; ++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    dfs1(1,0);
    dfs2(1,0,1);
    for(int i=1; i<=n; ++i) printf("%lld%c",ans[i],i==n?'\n':' ');
    return 0;
}

CF741D. Arpa’s letter-marked tree and Mehrdad’s Dokhtar-kosh paths

链接：https://codeforces.com/contest/741/problem/D

题意：给定一棵 n 个节点以 1 为根的树。边权为某个字符，输出每棵子树内最长回文路径的长度。（路径上的字符可以重新排列）

思路：一共有 22 个字符。回文路径可以重新拍，所以字符出现的次数为偶数次，且最多只有一个字符出现奇数次。那么就可以用 0 1 来表示奇偶。

这样就可以状压，用 01 串来表示当前的路径包含的字符。最后如果 dis[u] ^ dis[v] 的结果为 0 或者是 2 的幂，那么就表示组成了回文串。此时统计答案即可
用 d[i]来表示：路径值为 i 的最长长度。
步骤：先得到重儿子的信息，然后合并根节点，然后合并轻儿子，合并完一个轻儿子，就把信息更新上去（更新 d 数组），这样才能遍历完所有以经过 u 的路径
有一个细节：更新 ans 前需要判断 d 是否存在，否则可能因为 depth[v] 太大而更新错误，或者直接将 d 设为 -inf 。这样就取消了 depth[v] 的影响

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=5e5+5,maxm=(1<<23),inf=1e9;

int n;
vector<int> e[maxn];
int dfn[maxn],id[maxn],times;
int dis[maxm],d[maxm],ans[maxn];
int sz[maxn],depth[maxn],son[maxn];

void dfs1(int u,int fa)
{
    dfn[u]=++times;
    id[times]=u;
    sz[u]=1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        dis[v]^=dis[u];
        depth[v]=depth[u]+1;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
}
void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
        ans[u]=max(ans[u],ans[v]);
    }
    if(son[u]) dfs2(son[u],u,1),ans[u]=max(ans[u],ans[son[u]]);
    ans[u]=max(ans[u],d[dis[u]]-depth[u]);
    for(int i=1; i<=22; ++i) ans[u]=max(ans[u],d[dis[u]^(1<<(i-1))]-depth[u]);
    d[dis[u]]=max(depth[u],d[dis[u]]);
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        int l=dfn[v],r=dfn[v]+sz[v]-1;
        for(int i=l; i<=r; ++i)
        {
            int vv=id[i];
            ans[u]=max(ans[u],d[dis[vv]]+depth[vv]-2*depth[u]);
            for(int j=1; j<=22; ++j)
                ans[u]=max(ans[u],d[dis[vv]^(1<<j-1)]+depth[vv]-2*depth[u]);
        }
        for(int i=l; i<=r; ++i) d[dis[id[i]]]=max(d[dis[id[i]]],depth[id[i]]);
    }
    if(!kp)
    {
        int l=dfn[u],r=dfn[u]+sz[u]-1;
        for(int i=l; i<=r; ++i) d[dis[id[i]]]=-inf;
    }
}
int main()
{
    scanf("%d",&n);
    for(int v=2; v<=n; ++v)
    {
        int u;
        char c;
        scanf("%d %c",&u,&c);
        e[u].push_back(v);
        dis[v]=1<<(c-'a');
    }
    memset(d,-0x3f,sizeof(d));
    dfs1(1,0);
    dfs2(1,0,1);
    for(int i=1; i<=n; ++i) printf("%d%c",ans[i],i==n?'\n':' ');
    return 0;
}

CF375D. Tree and Queries

链接：https://codeforces.com/contest/375/problem/D

题意：给定一棵 n 个节点的树，根节点为 1。每个节点上有一个颜色 c_i。m 次操作。操作有一种：u k：询问在以 u 为根的子树中，出现次数 $\ge k$ 的颜色有多少种。 $(2\le n\le 10^5，1\le m\le 10^5，1\le c_i,k\le 10^5)$

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=1e5+5;

struct Opt
{
    int k,id;
};
int n,m;
int col[maxn],cnt[maxn],tot[maxn];
vector<int> e[maxn];
vector<Opt> opt[maxn];

int dfn[maxn],id[maxn],times;
int sz[maxn],son[maxn];
int ans[maxn];

void dfs1(int u,int fa)
{
    dfn[u]=++times;
    id[times]=u;
    sz[u]=1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
}
void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    tot[++cnt[col[u]]]++;
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        int l=dfn[v],r=dfn[v]+sz[v]-1;
        for(int i=l; i<=r; ++i)
        {
            int vv=id[i];
            tot[++cnt[col[vv]]]++;
        }
    }
    for(auto x: opt[u]) ans[x.id]=tot[x.k];
    if(!kp)
    {
        for(int i=1; i<=sz[u]; ++i) tot[i]=0;
        for(int i=dfn[u]; i<=dfn[u]+sz[u]-1; ++i) cnt[col[id[i]]]=0;
    }
}
int main()
{
    scanf("%d%d",&n,&m);
    for(int i=1; i<=n; ++i) scanf("%d",&col[i]);
    for(int i=1; i<=n-1; ++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    for(int i=1; i<=m; ++i)
    {
        int u,k;
        scanf("%d%d",&u,&k);
        opt[u].push_back({k,i});
    }
    dfs1(1,0);
    dfs2(1,0,1);
    for(int i=1; i<=m; ++i) printf("%d\n",ans[i]);
    return 0;
}

CF570D. Tree Requests

链接：https://codeforces.com/contest/570/problem/D

题意：给定一个以 1 为根的 n 个节点的树，每个点上有一个字母(a-z)，每个点的深度定义为该节点到1号节点路径上的点数.每次询问 a,b 。查询以 a 为根的子树内深度为 b 的节点上的字母重新排列之后是否能构成回文串. $(1\le n \le 5\times 10^5)$

思路：开一个数组 $c n t [i] [j]$ 表示深度为 $i$ 字符 $j$ 的出现次数。在每个子树暴力统计

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=5e5+5;

int n,m;
int cnt[maxn][30],val[maxn],ans[maxn];
char s[maxn];
struct Opt
{
    int d,id;
};
vector<int> e[maxn];
vector<Opt> opt[maxn];
int sz[maxn],son[maxn],depth[maxn];
int dfn[maxn],id[maxn],times;

bool check(int d)
{
    int res=0;
    for(int i=1; i<=26; ++i)
        res+=(cnt[d][i]&1);
    return res<=1;
}
void dfs1(int u,int fa)
{
    dfn[u]=++times;
    id[times]=u;
    sz[u]=1;
    depth[u]=depth[fa]+1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
}
void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    cnt[depth[u]][val[u]]++;
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        int l=dfn[v],r=dfn[v]+sz[v]-1;
        for(int i=l; i<=r; ++i)
        {
            int vv=id[i];
            cnt[depth[vv]][val[vv]]++;
        }
    }
    for(auto x: opt[u]) ans[x.id]=check(x.d);
    if(!kp)
    {
        int l=dfn[u],r=dfn[u]+sz[u]-1;
        for(int i=l; i<=r; ++i)
        {
            int vv=id[i];
            cnt[depth[vv]][val[vv]]=0;
        }
    }
}
int main()
{
    scanf("%d%d",&n,&m);
    for(int v=2; v<=n; ++v)
    {
        int u;
        scanf("%d",&u);
        e[u].push_back(v);
    }
    scanf("%s",s+1);
    for(int i=1; i<=n; ++i) val[i]=s[i]-'a'+1;
    for(int i=1; i<=m; ++i)
    {
        int u,d;
        scanf("%d%d",&u,&d);
        opt[u].push_back({d,i});
    }
    dfs1(1,0);
    dfs2(1,0,1);
    for(int i=1; i<=m; ++i) puts(ans[i]?"Yes":"No");
    return 0;
}

CF208E. Blood Cousins

链接：https://codeforces.com/contest/208/problem/E

题意：给定一片森林，每次询问一个点与多少个点 k 级祖先相同 $(1\le n ,m \le 10^5)$

思路：先倍增求 k 级祖先，然后在开一个桶，用 $c n t [i]$ 表示深度为 i 的点有多少个，这样就可以统计了，和上题差不多

细节：son[u] 需要初始化为 -1 。因为 0 也算节点。清空的时候，清空的是深度。

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=1e5+5;

int n,m;
vector<int> e[maxn];
int ans[maxn];
int sz[maxn],fa[maxn][22],son[maxn];
int in[maxn],out[maxn],id[maxn],times;
int cnt[maxn],depth[maxn];
struct Opt
{
    int k,id;
};
vector<Opt> opt[maxn];

void dfs1(int u,int f)
{
    in[u]=++times;
    id[times]=u;
    sz[u]=1;
    fa[u][0]=f;
    depth[u]=depth[f]+1;
    for(int i=1; i<=20; ++i)
        fa[u][i]=fa[fa[u][i-1]][i-1];
    for(auto v: e[u])
    {
        if(v==f) continue;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
    out[u]=times;
}
void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]!=-1) dfs2(son[u],u,1);
    cnt[depth[u]]++;
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        for(int i=in[v]; i<=out[v]; ++i)
        {
            int vv=id[i];
            cnt[depth[vv]]++;
        }
    }
    for(auto x: opt[u]) ans[x.id]=cnt[x.k+depth[u]]-1;
    if(!kp) for(int i=in[u]; i<=out[u]; ++i) cnt[depth[id[i]]]=0;
}
int main()
{
    scanf("%d",&n);
    for(int v=1; v<=n; ++v)
    {
        int u;
        scanf("%d",&u);
        e[u].push_back(v);
    }
    depth[0]=-1;
    memset(son,-1,sizeof(son));
    dfs1(0,0);
    scanf("%d",&m);
    for(int i=1; i<=m; ++i)
    {
        int u,k;
        scanf("%d%d",&u,&k);
        for(int i=0; (1<<i)<=k; ++i)
            if(k>>i&1) u=fa[u][i];
        if(u==0) ans[i]=0;
        else opt[u].push_back({k,i});
    }
    for(auto v: e[0]) dfs2(v,0,0);
    for(int i=1; i<=m; ++i) printf("%d%c",ans[i],i==m?'\n':' ');
    return 0;
}

CF246E. Blood Cousins Return

链接：https://codeforces.com/contest/246/problem/E

题意：给定一片森林。每个点有一个名字，每次询问一个点的 k 级儿子中，不同的名字个数。 $(1\le n ,m \le 10^5)$

思路：开一个 set 数组， $s e t [i]$ 表示深度为 $i$ 的不同颜色的个数

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=2e6+5;

int n,m;
vector<int> e[maxn];
string name[maxn];
unordered_map<string,int> vis;

int in[maxn],out[maxn],id[maxn],times;
int sz[maxn],ans[maxn],depth[maxn],son[maxn];

struct Opt
{
    int k,id;
};
vector<Opt> opt[maxn];
int col[maxn];
set<int> s[maxn];

void dfs1(int u,int fa)
{
    in[u]=++times;
    id[times]=u;
    sz[u]=1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        depth[v]=depth[u]+1;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
    out[u]=times;
}

void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    s[depth[u]].insert(col[u]);
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        for(int i=in[v]; i<=out[v]; ++i)
        {
            int vv=id[i];
            s[depth[vv]].insert(col[vv]);
        }
    }
    for(auto x: opt[u]) ans[x.id]=s[x.k+depth[u]].size();
    if(!kp) for(int i=in[u]; i<=out[u]; ++i) s[depth[id[i]]].clear();
}

int main()
{
    ios::sync_with_stdio(0);
    cin>>n;
    int color=0;
    for(int v=1; v<=n; ++v)
    {
        int u;
        cin>>name[v]>>u;
        e[u].push_back(v);
        if(!vis[name[v]]) vis[name[v]]=++color;
        col[v]=vis[name[v]];
    }
    for(auto v: e[0]) dfs1(v,0);
    cin>>m;
    for(int i=1; i<=m; ++i)
    {
        int u,k;
        cin>>u>>k;
        opt[u].push_back({k,i});
    }
    for(auto v: e[0]) dfs2(v,0,0);
    for(int i=1; i<=m; ++i) cout<<ans[i]<<"\n";
    return 0;
}

CF1009F. Dominant Indices

链接：https://codeforces.com/contest/1009/problem/F

题意：给定一棵以 $1$ 为根， $n$ 个节点的树。设 $d (u, x)$ 为 $u$ 子树中到 $u$ 距离为 $x$ 的节点数。对于每个点，求一个最小的 $k$ ，使得 $d (u, k)$ 最大。

思路：

dsu on tree ：一个桶维护一个每个深度的点数，然后每次更新时取最大就好了。
线段树合并：按照深度开一棵权值线段树，每次查询最左边的最大值即可

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=1e6+5;

int n;
vector<int> e[maxn];
int in[maxn],out[maxn],sz[maxn],id[maxn],times;
int depth[maxn],ans[maxn],son[maxn];

void dfs1(int u,int fa)
{
    in[u]=++times;
    id[times]=u;
    sz[u]=1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        depth[v]=depth[u]+1;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
    out[u]=times;
}
int k,mx,cnt[maxn];
void update(int u)
{
    cnt[depth[u]]++;
    if(cnt[depth[u]]>mx) mx=cnt[depth[u]],k=depth[u];
    else if(cnt[depth[u]]==mx) if(depth[u]<k) k=depth[u];
}

void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    update(u);
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        for(int i=in[v]; i<=out[v]; ++i) update(id[i]);
    }
    ans[u]=k;
    if(!kp)
    {
        for(int i=in[u]; i<=out[u]; ++i) cnt[depth[id[i]]]=0;
        mx=0,k=0;
    }
}

int main()
{
    scanf("%d",&n);
    for(int i=1; i<=n-1; ++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    dfs1(1,0);
    dfs2(1,0,0);
    for(int i=1; i<=n; ++i) printf("%d\n",ans[i]-depth[i]);
    return 0;
}

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=1e6+5,tot=1e6;

int n;
vector<int> e[maxn];
int depth[maxn],ans[maxn];

int root[maxn],ls[maxn*25],rs[maxn*25],st[maxn*25],no;

void update(int &rt,int pos,int L,int R)
{
    if(!rt) rt=++no;
    if(L==R)
    {
        st[rt]++;
        return;
    }
    int mid=(L+R)>>1;
    if(pos<=mid) update(ls[rt],pos,L,mid);
    else update(rs[rt],pos,mid+1,R);
    st[rt]=max(st[ls[rt]],st[rs[rt]]);
}
int query(int rt,int L,int R)
{
    if(!rt) return 0;
    if(L==R) return L;
    int mid=(L+R)>>1;
    if(st[ls[rt]]>=st[rs[rt]]) return query(ls[rt],L,mid);
    else return query(rs[rt],mid+1,R);
}
int merge(int rt1,int rt2,int L,int R)
{
    if(!rt1||!rt2) return rt1?rt1:rt2;
    int rt=rt1;
    if(L==R)
    {
        st[rt]+=st[rt2];
        return rt;
    }
    int mid=(L+R)>>1;
    ls[rt]=merge(ls[rt1],ls[rt2],L,mid);
    rs[rt]=merge(rs[rt1],rs[rt2],mid+1,R);
    st[rt]=max(st[ls[rt]],st[rs[rt]]);
    return rt;
}
void dfs(int u,int fa)
{
    depth[u]=depth[fa]+1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        dfs(v,u);
        root[u]=merge(root[u],root[v],1,tot);
    }
    update(root[u],depth[u],1,tot);
    ans[u]=query(root[u],1,tot);
}
int main()
{
    scanf("%d",&n);
    for(int i=1; i<=n-1; ++i)
    {
        int u,v;
        scanf("%d%d",&u,&v);
        e[u].push_back(v);
        e[v].push_back(u);
    }
    dfs(1,0);
    for(int i=1; i<=n; ++i)
        printf("%d\n",ans[i]-depth[i]);
    return 0;
}

2019 ICPC Nanchang K. Tree

链接：https://www.jisuanke.com/contest/5530/challenges

题意：给定一棵 n 个节点以 1 为根树，每个点有点权 $v_i$ ，然后询问有序对（x，y）满足一下条件：

x 不是 y 的祖先，且 y 不是 x 的祖先
x 和 y 的简单路径之和不超过 k
x 和 y 的点权满足： $v_x+v_y=2v_{lca(x,y)}$

思路：假设当前枚举的点为 $v_1$ ，子树根节点为 u 。需要找的点 $v_2$ 需要满足的条件是： $val_{v_2}=2val_{u}-val_{v_1}$ ， $1\le depth_{v_2}\le k +2depth_u-depth_{v_1}$ ，大于等于 1 是要满足不互为祖先的条件。

可以看到 val 是固定的，因此可以对每一个 val 建一棵权值线段树，然后根据 depth 的范围查询相应的点对。注意要查询一棵子树，然后更新一棵子树。这样可以保证枚举的 lca 的所有路径都是查询到。
求的是有序对，最后答案 $\times 2$

#include <bits/stdc++.h>
#define ll long long
using namespace std;
const int maxn=1e5+5,N=1e5;
int n,k;
vector<int> e[maxn];
ll ans;
int in[maxn],out[maxn],id[maxn],times;
int sz[maxn],depth[maxn],son[maxn],val[maxn];
int root[maxn],st[maxn*100],ls[maxn*100],rs[maxn*100],no;
void update(int &rt,int pos,int L,int R,int val)
{
    if(!rt) rt=++no;
    st[rt]+=val;
    if(L==R) return;
    int mid=(L+R)>>1;
    if(pos<=mid) update(ls[rt],pos,L,mid,val);
    if(pos>mid) update(rs[rt],pos,mid+1,R,val);
}
int query(int rt,int l,int r,int L,int R)
{
    if(!rt) return 0;
    if(l<=L&&R<=r) return st[rt];
    int mid=(L+R)>>1;
    int ans=0;
    if(l<=mid) ans+=query(ls[rt],l,r,L,mid);
    if(r>mid) ans+=query(rs[rt],l,r,mid+1,R);
    return ans;
}
void dfs1(int u,int fa)
{
    in[u]=++times;
    id[times]=u;
    sz[u]=1;
    for(auto v: e[u])
    {
        if(v==fa) continue;
        depth[v]=depth[u]+1;
        dfs1(v,u);
        sz[u]+=sz[v];
        if(sz[v]>sz[son[u]]) son[u]=v;
    }
    out[u]=times;
}
void dfs2(int u,int fa,int kp)
{
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u,0);
    }
    if(son[u]) dfs2(son[u],u,1);
    for(auto v: e[u])
    {
        if(v==fa||v==son[u]) continue;
        for(int i=in[v]; i<=out[v]; ++i)
        {
            int vv=id[i];
            int rt=2*val[u]-val[vv];
            int r=k+2*depth[u]-depth[vv];
            if(rt>=0&&rt<=100000&&r>=1) ans+=query(root[rt],0,r,0,N);
        }
        for(int i=in[v]; i<=out[v]; ++i)
        {
            int vv=id[i];
            update(root[val[vv]],depth[vv],0,N,1);
        }
    }
    update(root[val[u]],depth[u],0,N,1);
    if(!kp)
    {
        for(int i=in[u]; i<=out[u]; ++i)
        {
            int vv=id[i];
            update(root[val[vv]],depth[vv],0,N,-1);
        }
    }
}
int main()
{
    scanf("%d%d",&n,&k);
    for(int i=1; i<=n; ++i) scanf("%d",&val[i]);
    for(int v=2; v<=n; ++v)
    {
        int u;
        scanf("%d",&u);
        e[u].push_back(v);
    }
    dfs1(1,0);
    dfs2(1,0,1);
    printf("%lld\n",2ll*ans);
    return 0;
}

原文链接：https://blog.csdn.net/cheng__yu_/article/details/109298143