代码-洛谷P1600 天天爱跑步

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef double db;
#define x first
#define y second
#define bg begin()
#define ed end()
#define pb push_back
#define mp make_pair
#define sz(a) int((a).size())
#define R(i,n) for(int i(0);i<(n);++i)
#define L(i,n) for(int i((n)-1);~i;--i)
const int iinf=0x3f3f3f3f;
const ll linf=0x3f3f3f3f3f3f3f3f;

//Data
const int N=299998;
int n,m,wt[N],ns[N];

//SegmentTree
const int tN=1e7;
#define mid ((l+r)>>1)
int tn,ls[tN],rs[tN],val[tN];
int newt(){return ls[tn]=rs[tn]=-1,val[tn]=0,tn++;}
void add(int i,int v,int &u,int l=0,int r=n<<1){
    if(r<=i||i+1<=l) return; if(!~u) u=newt();
    if(r-l==1) return /*cout<<"["<<l<<"]+="<<v<<'\n',*/val[u]+=v,void();
    add(i,v,ls[u],l,mid),add(i,v,rs[u],mid,r);
}
int get(int i,int u,int l=0,int r=n<<1){
    if(r<=i||i+1<=l) return 0; if(!~u) return 0;
    if(r-l==1) return /*cout<<"I need ["<<l<<"]\n",*/val[u];
    return get(i,ls[u],l,mid)+get(i,rs[u],mid,r);
}
int merge(int u,int v,int l=0,int r=n<<1){
    if(!~u||!~v) return u+v+1;
    if(r-l==1) return val[u]+=val[v],u;
    ls[u]=merge(ls[u],ls[v],l,mid),rs[u]=merge(rs[u],rs[v],mid,r);
    return u;
}

//Tree
const int LG=20;
vector<int> e[N];
int dep[N],fa[LG][N];
void dfsi(int u){
    R(i,LG-1) ~fa[i][u]&&(fa[i+1][u]=fa[i][fa[i][u]]);
    for(int v:e[u]) (v^fa[0][u])&&(dep[v]=dep[fa[0][v]=u]+1,dfsi(v),true);
}
int lca(int u,int v){
    dep[u]<dep[v]&&(swap(u,v),true);
    L(i,LG) ~fa[i][u]&&dep[fa[i][u]]>=dep[v]&&(u=fa[i][u]); if(u==v) return u;
    L(i,LG) ~fa[i][u]&&~fa[i][v]&&fa[i][u]!=fa[i][v]&&(u=fa[i][u],v=fa[i][v]);
    return fa[0][u];
}
int ar[N],br[N];
void dfss(int u){
    for(int v:e[u])if(v^fa[0][u]) dfss(v),
        ar[u]=merge(ar[u],ar[v]),br[u]=merge(br[u],br[v]);
    // cout<<"u="<<u<<'\n';
    ns[u]=get(dep[u]+wt[u],ar[u])+get(dep[u]-wt[u]+n,br[u]);
    // cout<<get(dep[u]+wt[u],ar[u])<<" "<<get(dep[u]-wt[u]+n,br[u])<<'\n';
}

//Main
int main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>n>>m;
    R(i,n-1){int u,v; cin>>u>>v,--u,--v,e[u].pb(v),e[v].pb(u);}
    R(i,LG)R(u,n) fa[i][u]=-1; dfsi(0);
    R(u,n) cin>>wt[u],ar[u]=br[u]=-1;
    R(i,m){
        int u,v; cin>>u>>v,--u,--v; int a=lca(u,v),f=fa[0][a];
        // cout<<u<<" "<<v<<":"<<a<<" "<<f<<'\n';
        add(dep[u],1,ar[u]),~f&&(add(dep[u],-1,ar[f]),true);
        add(n+2*dep[a]-dep[u],1,br[v]),add(n+2*dep[a]-dep[u],-1,br[a]);
        // cout<<dep[u]<<" "<<n+2*dep[a]-dep[u]<<'\n';

    }
    dfss(0); R(i,n) cout<<ns[i]<<" ";cout<<'\n';
    return 0;
}