大小接近的点对（主席树）

题目描述

一天，Chika 对大小接近的点对产生了兴趣，她想搞明白这个问题的树上版本，你能帮助她吗？Chika 会给 你一棵有根树，这棵树有 n 个结点，被编号为 1 n，1 号结点是根。每个点有一个权值，i 号结点的权值为 a[i]。如果 u 是 v 的祖先结点，并且 abs(a[u]−a[v]) ≤K，那么 (u,v) 被称作一个“** 大小接近的点对 **”。 对于树上的每个结点 i，你都需要计算以其为根的子树中的“大小接近的点对”的数量。你需要知道：  
(1) abs(x) 代表 x 的绝对值。  
(2) 每个结点都是其自身的祖先结点. 

 

 

输入

输入文件的第一行包含两个整数 n (1≤n≤105) 和 k (1≤k≤109)，代表树中结点总数, 以及“大小接近的点对”的大小之差的上界。  
第二行包含 n 个整数，第 i 个整数是 a[i] (1≤ a[i] ≤109)，代表 i 号结点的权值。  
第三行包含 n−1 个整数，第 i 个整数是 i+1 号结点的父结点。 
 

输出

输出应该包含n行，每一行包括一个整数。第i行的整数代表以i为根的子树中的“大小接近的点对”的数量。

样例输入 Copy

7 5
2 4 4 1 4 6 4
1 2 3 1 2 3

样例输出 Copy

19
11
5
1
1
1
1

#include<iostream>
#include<queue>
#include<cstdio>
#include<cstring>
#include<iostream>
#include<string>
#include<algorithm>
#include<vector>
#include<cmath>
#include<stack>
#include<bitset>
using namespace std;
#define N 100010
#define LL long long

struct Node
{
    int l,r,data;
} T[N*20];

int first[N],nex[N],v[N],Ls[N],Rs[N],b[N*3],root[N];
int n,K,num,top,cnt,ans;
LL sum[N];

struct data
{
    int p,v;
} t[N*3];
bool cmp(data x,data y)
{
    return x.v<y.v;
}

void build(int pre,int cur,int p,int l,int r)//建树
{
    if(l==r)//如果是叶子节点，data为1返回
    {
        T[cur].data=T[pre].data+1;
        return;
    }
    int mid=(l+r)/2;//取区间中值进行比较
    if(p<=mid)//小于向左树查找
    {
        T[cur].l=++ans;
        T[cur].r=T[pre].r;
        build(T[pre].l,T[cur].l,p,l,mid);//重新建立一个新树
    }
    else
    {
        T[cur].r=++ans;
        T[cur].l=T[pre].l;
        build(T[pre].r,T[cur].r,p,mid+1,r);//重新建立一个新树
    }
    T[cur].data=T[T[cur].l].data+T[T[cur].r].data;//统计当前节点的元素总合
}
int Query(int L,int R,int ql,int qr,int l,int r)//区间查询
{
    if(ql<=l&&r<=qr)//当查询的区间正好在同一侧，返回
    {
        return T[R].data-T[L].data;
    }
    int mid=(l+r)/2;
    int res=0;//否则统计左右区间的和
    if(ql<=mid)
    {
        res+=Query(T[L].l,T[R].l,ql,qr,l,mid);
    }
    if(mid<qr)
    {
        res+=Query(T[L].r,T[R].r,ql,qr,mid+1,r);
    }
    return res;
}

void dfs(int x)
{
    Ls[x]=++top;
    root[Ls[x]]=++ans;
    build(root[Ls[x]-1],root[Ls[x]],b[x*3-2],1,cnt);//建树
    for(int i=first[x]; i; i=nex[i])//邻接表递归下一个点
    {
        dfs(v[i]);
        sum[x]+=sum[v[i]];
    }
    Rs[x]=top;
    sum[x]+=(LL)Query(root[Ls[x]-1],root[Rs[x]],b[x*3-1],b[x*3],1,cnt);
}
int main()
{
    int x;
    cnt=0,top=0,num=0;
    scanf("%d%d",&n,&K);
    for(int i=1; i<=n; ++i)//离散化处理
    {
        scanf("%d",&t[i*3-2].v);
        t[i*3-2].p=i*3-2;
        t[i*3-1].v=t[i*3-2].v-K;
        t[i*3-1].p=i*3-1;
        t[i*3].v=t[i*3-2].v+K;
        t[i*3].p=i*3;
    }
    sort(t+1,t+n*3+1,cmp);
    b[t[1].p]=++cnt;
    for(int i=2; i<=n*3; ++i)
    {
        if(t[i].v!=t[i-1].v)
        {
            ++cnt;
        }
        b[t[i].p]=cnt;
    }
    for(int i=2; i<=n; ++i)
    {
        scanf("%d",&x);
        v[++num]=i;
        nex[num]=first[x];
        first[x]=num;
    }
    dfs(1);
    for(int i=1; i<=n; ++i)
    {
        printf("%lld\n",sum[i]);
    }
    return 0;
}