codeforces 293E Close Vertices 点分治+滑窗+treap

http://codeforces.com/contest/293/problem/E

题意：求树上合法点对的个数。合法条件为：路径长度<=W，路径边数<=L。

显然是点分治。求解的时候第一维直接滑窗，第二维插入treap中统计，滑窗的时候删掉即可。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<vector>
#include<set>
#define REP(i,a,b) for(int i=a;i<=b;i++)
#define MS0(a) memset(a,0,sizeof(a))

using namespace std;

typedef long long ll;
const int maxn=1000100;
const int INF=1e9+10;

int N;
ll L,W;
int p;ll w;
struct Edge
{
    int v;ll w;
};Edge e[maxn*2];int tot;
int first[maxn],Next[maxn*2];
bool vis[maxn];
struct PLL
{
    ll l,w;
    friend bool operator<(PLL A,PLL B)
    {
        return A.l==B.l?A.w<B.w:A.l<B.l;
    }
};PLL d[maxn];int dn;
int rt,balance;

void Init()
{
    tot=0;
    memset(first,-1,sizeof(first));
}

void addedge(int u,int v,ll w)
{
    e[++tot]={v,w};
    Next[tot]=first[u];
    first[u]=tot;
}

int get_rt(int u,int f,int sz)
{
    int cnt=1,balance1=0;
    for(int i=first[u];~i;i=Next[i]){
        int v=e[i].v;
        if(v==f||vis[v]) continue;
        int tmp=get_rt(v,u,sz);
        cnt+=tmp;
        balance1=max(balance1,tmp);
    }
    balance1=max(balance1,sz-cnt);
    if(balance1<balance){
        rt=u;
        balance=balance1;
    }
    return cnt;
}

void dfs_d(int u,int f,ll len,ll sw)
{
    d[++dn]={len,sw};
    for(int i=first[u];~i;i=Next[i]){
        int v=e[i].v;ll w=e[i].w;
        if(v==f||vis[v]) continue;
        dfs_d(v,u,len+1,sw+w);
    }
}

///---Treap
struct Node
{
    int ch[2];
    int v;
    int r;
    int sz;
};Node tr[4*maxn];int tot1;
int s[maxn],tot2;
int root;

int newnode()
{
    int k;
    if(tot2) k=s[tot2--];
    else k=++tot1;
    tr[k].ch[0]=tr[k].ch[1]=-1;
    tr[k].v=0;
    tr[k].r=rand();
    tr[k].sz=1;
    return k;
}

void up(int o)
{
    tr[o].sz=1;
    if(~tr[o].ch[0]) tr[o].sz+=tr[tr[o].ch[0]].sz;
    if(~tr[o].ch[1]) tr[o].sz+=tr[tr[o].ch[1]].sz;
}

void rot(int &o,int d)
{
    int k=tr[o].ch[d^1];
    tr[o].ch[d^1]=tr[k].ch[d];
    tr[k].ch[d]=o;
    up(o);up(k);
    o=k;
}

void Insert(int &o,ll x)
{
    if(o==-1){
        o=newnode();
        tr[o].v=x;
        return;
    }
    int d=x<tr[o].v?0:1;
    Insert(tr[o].ch[d],x);
    if(tr[tr[o].ch[d]].r>tr[o].r) rot(o,d^1);
    up(o);
}

void Del(int &o,ll x)
{
    if(o==-1) return;
    if(tr[o].v==x){
        if(tr[o].ch[0]==-1) s[++tot2]=o,o=tr[o].ch[1];
        else if(tr[o].ch[1]==-1) s[++tot2]=o,o=tr[o].ch[0];
        else{
            int d=tr[tr[o].ch[0]].r>tr[tr[o].ch[1]].r?0:1;
            rot(o,d^1);
            Del(tr[o].ch[d^1],x);
        }
        if(~o) up(o);
        return;
    }
    if(x<tr[o].v) Del(tr[o].ch[0],x);
    else Del(tr[o].ch[1],x);
    up(o);
}

ll query(int o,ll x)
{
    if(o==-1) return 0;
    if(x<tr[o].v) return query(tr[o].ch[0],x);
    ll ls=tr[o].ch[0]==-1?0:tr[tr[o].ch[0]].sz;
    return ls+1+query(tr[o].ch[1],x);
}

ll calc(int u,int f,ll len,ll sw)
{
    dn=0;
    dfs_d(u,f,len,sw);
    sort(d+1,d+dn+1);
    ll res=0;
    root=-1;tot1=tot2=0;
    REP(i,1,dn) Insert(root,d[i].w);
    for(int l=1,r=dn;l<r;l++){
        while(r>l&&d[l].l+d[r].l>L) Del(root,d[r].w),r--;
        Del(root,d[l].w);
        res+=query(root,W-d[l].w);

    }
    return res;
}

ll solve(int u)
{
    rt=u,balance=INF;
    int sz=get_rt(u,0,N);
    rt=u,balance=INF;
    get_rt(u,0,sz);
    u=rt;
    vis[u]=1;
    ll res=0;
    res+=calc(u,0,0,0);
    for(int i=first[u];~i;i=Next[i]){
        int v=e[i].v;ll w=e[i].w;
        if(vis[v]) continue;
        res-=calc(v,u,1,w);
    }
    for(int i=first[u];~i;i=Next[i]){
        int v=e[i].v;ll w=e[i].w;
        if(vis[v]) continue;
        res+=solve(v);
    }
    return res;
}

int main()
{
    //test();return 0;
    freopen("in.txt","r",stdin);
    while(cin>>N>>L>>W){
        Init();
        REP(i,1,N-1){
            scanf("%d%I64d",&p,&w);
            addedge(p,i+1,w);
            addedge(i+1,p,w);
        }
        MS0(vis);
        printf("%I64d
",solve(1));
    }
    return 0;
}