2019 ICPC 南昌 Regional K. Tree（树上启发式+动态开点线段树）

传送门

给定了一棵带点权的有根树，给定正整数k，求出节点对(i,j)的个数，满足a[i]+a[j]=2*a[lca]，且dis(i,j)<=k，且i j没有祖先关系。

---

通过LCA，将限制条件转化成如下式子：
dis(i,j)=dis(1,i)+dis(1,j) - 2*dis(1,lca) <= k

a[i]+a[j]=2*a[lca]

我们枚举LCA计算以它为根的树的答案，
对于其内顶点i,我们要在其兄弟子树中找到点权为2*a[lca]-a[i],且满足dis(j)<=k+2*dis(lca)-dis(i)的节点个数。

后面的这个不等式其实是一个区间查询，我们对每个不同的点权都建一个线段树，定义域是dep[i]，维护对应深度区间内，这个点权出现了多少次， 于是每次查询区间和即可 。

由于每个点权都对应了一棵线段树，所以需要动态开点。

做法大概是：对于rt，我们一棵一棵地遍历它的子树，每棵子树都是先统计答案，然后更新线段树的信息。 如果我们枚举每棵子树，然后清空信息，这样是标准的O(N²)， 但是我们发现，最多可以保留1棵子树的信息不清空，因为回到rt时这颗子树和0棵子树的信息不会产生答案贡献，从第二颗子树开始才会产生贡献， 所以我们希望让保留信息的子树尽量大，于是结合重儿子的概念，保留重子树，这样可以证明是O(Nlogn)

#include<bits/stdc++.h>
using namespace std;
//#pragma GCC optimize(2)
#define ull unsigned long long
#define ll long long
#define pii pair<int, int>
const int maxn = 4e5 + 10;
const ll mod = 998244353;
const ll inf = (ll)4e16+5;
const int INF = 1e9 + 7;
const double pi = acos(-1.0);
ll inv(ll b){if(b==1)return 1;return(mod-mod/b)*inv(mod%b)%mod;}
inline ll read()
{
    ll x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
//给定带点权的有根树 给定整数k 求满足最短距离小于等于k a[i]+a[j]==2*a[lca] 的(i,j)对的个数
//dsu on tree 枚举LCA 对于点a[i] 找到其他相邻子树中 
//满足a[j]=2*a[rt]-a[i] dep[j]<=k+2dep[rt]-dep[i]的节点个数 带约束性的权值个数问题 
//由于点权是确定的 每个点权 对深度建树 深度为i 有多少个这样的点权 维护区间和
//动态开点
int root[maxn],lc[maxn*40],rc[maxn*40],tot,tree[maxn*40]; 
int n,k;
vector<int> g[maxn];
int a[maxn],fa[maxn];
int sz[maxn],son[maxn],in[maxn],clk,pos[maxn];
int dep[maxn];
ll ans=0;
void dfs(int rt)
{
    sz[rt]=1;
    in[rt]=++clk;
    pos[clk]=rt;
    dep[rt]=dep[fa[rt]]+1;
    for(int i:g[rt]) 
    {
        dfs(i);
        sz[rt]+=sz[i];
        if(sz[i] > sz[son[rt]]) son[rt]=i;
    }
}
inline void pushup(int rt)
{
    tree[rt]=tree[lc[rt]] + tree[rc[rt]];
}
inline void upd(int &rt,int l,int r,int pos,int v)
{
    if(!rt) rt=++tot;
    if(l==r) 
    {
        tree[rt]+=v;
        return ;
    }
    int mid=l+r>>1;
    if(pos<=mid) upd(lc[rt],l,mid,pos,v);
    else upd(rc[rt],mid+1,r,pos,v);
    pushup(rt);
}
inline int qry(int rt,int l,int r,int vl,int vr)
{
    if(r<vl || l>vr || !rt) return 0;
    if(vl<=l && r<=vr) return tree[rt];
    int mid=l+r>>1;
    return qry(lc[rt],l,mid,vl,vr) + qry(rc[rt],mid+1,r,vl,vr);
}
int lca;
inline void get(int rt)
{
    for(int i=in[rt];i<in[rt]+sz[rt];i++)
    {
        int u=pos[i];
        int v=2*a[lca]-a[u];
        int R=k+2*dep[lca]-dep[u];
        ans+=qry(root[v],1,n,1,R);
    }
}
inline void add(int rt,int v)
{
    for(int i=in[rt];i<in[rt]+sz[rt];i++)
    {
        int u=pos[i];
        upd(root[a[u]],1,n,dep[u],v);
    }
}
void dfs2(int rt,bool save)
{
    for(int i:g[rt]) 
    {
        if(i==son[rt]) continue;
        dfs2(i,0);
    }
    if(son[rt]) dfs2(son[rt],1);
    //lca 不会计入答案 枚举每颗子树 先统计答案再更新
    lca=rt;
    for(int i:g[rt]) 
    {
        if(i==son[rt]) continue;
        get(i);
        add(i,1);
    }
    upd(root[a[rt]],1,n,dep[rt],1);
    if(!save)
    {
        add(rt,-1);
    }
}
int main()
{
    scanf("%d %d",&n,&k);
    for(int i=1;i<=n;i++) scanf("%d",a+i);
    for(int i=2;i<=n;i++)
    {
        scanf("%d",fa+i);
        g[fa[i]].push_back(i);
    }
    dfs(1);
    dfs2(1,1);
    cout<<2ll*ans<<'\n';
    return 0;
}