luoguP1600 天天爱跑步(NOIP2016)(主席树+树链剖分)

luogu P1600 天天爱跑步 题目

首先声明一下：这篇的代码是tyher巨佬的，不是我的……方法也是他的原创

#include<iostream>
#include<cstdlib>
#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
#include<vector>
#include<queue>
#define il inline
#define rg register
#define ll long long
#define ld long double
#define N 300000
#define inf 2147483647
using namespace std;
int n,m,u,v,cnt,num;
int hd[N],w[N],s[N],t[N];
int de[N],top[N],son[N],tot[N];
int fa[N],idx[N];
struct Edge{
    int nt,to;
}edge[N<<1];
int rt1[N*2],rt2[N*2],ls[N*40],rs[N*40],val[N*40];
il void re(rg int &x);
il void link(rg int fm,rg int to);
void Dfs1(rg int i,rg int fm);
void Dfs2(rg int i,rg int head);
int lca(rg int x,rg int y);
il void work(rg int s,rg int t);
void add(rg int now,rg int le,rg int ri,rg int pos);
void update(rg int now,rg int le,rg int ri,rg int L,rg int R,rg int w);
int query(rg int now,rg int le,rg int ri,rg int pos);
int main(){
    re(n),re(m);
    for(rg int i=1;i<n;++i){
        re(u),re(v);
        link(u,v),link(v,u);
    }
    for(rg int i=1;i<=n;++i)
        re(w[i]);
    for(rg int i=1;i<=m;++i)
        re(s[i]),re(t[i]);
    cnt=0,Dfs1(1,0),Dfs2(1,1);
    for(rg int i=1;i<=m;++i)
        work(s[i],t[i]);
    for(rg int i=1;i<=n;++i){
        int ans1=query(rt1[de[i]+w[i]],1,n,idx[i]);
        int ans2=query(rt2[w[i]-de[i]+N],1,n,idx[i]);
        printf("%d ",ans1+ans2);
    }
    return 0;
}
int query(rg int now,rg int le,rg int ri,rg int pos){
    if(!now)return 0;
    if(le==ri)return val[now];
    if(val[now]){
        if(ls[now])val[ls[now]]+=val[now];
        if(rs[now])val[rs[now]]+=val[now];
        val[now]=0;
    }
    rg int mid=((le+ri)>>1);
    if(pos<=mid)return query(ls[now],le,mid,pos);
    else return query(rs[now],mid+1,ri,pos);
}
il void work(rg int s,rg int t){
    rg int Lca=lca(s,t),u=s;
    while(top[u]!=top[Lca]){
        update(rt1[de[s]],1,n,idx[top[u]],idx[u],1);
        u=fa[top[u]];
    }
    update(rt1[de[s]],1,n,idx[Lca],idx[u],1);
    u=t;
    while(top[u]!=top[Lca]){
        update(rt2[de[s]-2*de[Lca]+N],1,n,idx[top[u]],idx[u],1);
        u=fa[top[u]];
    }
    update(rt2[de[s]-2*de[Lca]+N],1,n,idx[Lca],idx[u],1);
    update(rt1[de[s]],1,n,idx[Lca],idx[Lca],-1);
}
void add(rg int now,rg int le,rg int ri,rg int pos){
    if(le==ri)return;
    rg int mid=((le+ri)>>1);
    if(pos<=mid){
        if(!ls[now])ls[now]=(++num);
        add(ls[now],le,mid,pos);
    }
    else{
        if(!rs[now])rs[now]=(++num);
        add(rs[now],mid+1,ri,pos);
    }
}
void update(rg int now,rg int le,rg int ri,rg int L,rg int R,rg int w){
    if(!now)return;
    if(L==le&&R==ri){val[now]+=w;return;}
    rg int mid=((le+ri)>>1);
    if(R<=mid)update(ls[now],le,mid,L,R,w);
    else if(L>mid)update(rs[now],mid+1,ri,L,R,w);
    else update(ls[now],le,mid,L,mid,w),update(rs[now],mid+1,ri,mid+1,R,w);
}
int lca(rg int x,rg int y){
    while(top[x]!=top[y]){
        if(de[top[x]]<de[top[y]])swap(x,y);
        x=fa[top[x]];
    }
    if(de[x]>de[y])return y;
    else return x;
}   
void Dfs1(rg int i,rg int fm){
    de[i]=de[fm]+1,fa[i]=fm;
    tot[i]=1;
    rg int maxn=0;
    for(rg int k=hd[i];k;k=edge[k].nt){
        rg int qw=edge[k].to;
        if(qw==fm)continue;
        Dfs1(qw,i),tot[i]+=tot[qw];
        if(tot[qw]>maxn)maxn=tot[qw],son[i]=qw;
    }
}
void Dfs2(rg int i,rg int head){
    idx[i]=(++cnt),top[i]=head;
    if(!rt1[de[i]+w[i]])rt1[de[i]+w[i]]=(++num);
    add(rt1[de[i]+w[i]],1,n,idx[i]);
    if(!rt2[w[i]-de[i]+N])rt2[w[i]-de[i]+N]=(++num);
    add(rt2[w[i]-de[i]+N],1,n,idx[i]);
    if(!son[i])return;
    Dfs2(son[i],head);
    for(rg int k=hd[i];k;k=edge[k].nt)
        if(!idx[edge[k].to])
            Dfs2(edge[k].to,edge[k].to);
}
il void re(rg int &x){
    rg int res=0;rg int w=1;char c=getchar();
    while((c<'0'||c>'9')&&c!='-')c=getchar();
    if(c=='-')w=-1,c=getchar();
    while(c>='0'&&c<='9')res=(res<<3)+(res<<1)+c-'0',c=getchar();
    x=w*res;
}
il void link(rg int fm,rg int to){
    edge[++cnt].nt=hd[fm];
    edge[cnt].to=to;
    hd[fm]=cnt;
}