bzoj4719: [Noip2016]天天爱跑步

题目
题解

#include<bits/stdc++.h>
using namespace std;
const int N=300003,M=600003;
struct kk{
    int u,v,lca,dis;
}p[N];
struct node{
    int to,ne;
}e[M];
int b[N+M],cnt[M],ans[N],w[N],dis[N],f[N][20],tot,i,u,v,h[N],n,m,vis[N];
vector<int>v1[M],v2[M],v3[M];
int read(){
    int x=0;char c;
    do c=getchar();while (c<'0' || c>'9');
    while ('0'<=c && c<='9') x=(x<<3)+(x<<1)+(c^48),c=getchar();
    return x;
}
void add(int u,int v){
    e[++tot]=(node){v,h[u]};
    h[u]=tot;
}
void dfs(int u,int dep){
    vis[u]=1;
    for (int i=1;i<20;i++) f[u][i]=f[f[u][i-1]][i-1];
    for (int i=h[u],v;i;i=e[i].ne)
        if (!vis[v=e[i].to]){
            dis[v]=dis[u]+1;
            f[v][0]=u;
            dfs(v,dep+1);
        }
    vis[u]=0;
}
int lca(int u,int v){
    if (dis[u]<dis[v]) swap(u,v);
    for (int i=0;i<20;i++)
        if ((dis[u]-dis[v])&(1<<i)) u=f[u][i];
    for (int i=19;i>=0;i--)
        if (f[u][i]!=f[v][i]) u=f[u][i],v=f[v][i];
    return u==v?u:f[u][0];
}
void dfs1(int u){
    vis[u]=1;
    int pre=b[dis[u]+w[u]+N];
    for (int i=h[u],v;i;i=e[i].ne)
        if (!vis[v=e[i].to]) dfs1(v);
    b[dis[u]+N]+=cnt[u];
    ans[u]+=b[dis[u]+w[u]+N]-pre;
    for (int i=0;i<v1[u].size();i++) b[dis[v1[u][i]]+N]--;
    vis[u]=0;
}
void dfs2(int u){
    vis[u]=1;
    int pre=b[w[u]-dis[u]+N];
    for (int i=h[u],v;i;i=e[i].ne)
        if (!vis[v=e[i].to]) dfs2(v);
    for (int i=0;i<v2[u].size();i++) b[v2[u][i]+N]++;
    ans[u]+=b[w[u]-dis[u]+N]-pre;
    for (int i=0;i<v3[u].size();i++) b[v3[u][i]+N]--;
    vis[u]=0;
}
int main(){
    n=read();m=read();
    for (i=1;i<n;i++) u=read(),v=read(),add(u,v),add(v,u);
    for (i=1;i<=n;i++) w[i]=read();
    f[1][0]=1;
    dfs(1,0);//预处理f数组
    for (i=1;i<=m;i++){
        p[i].u=u=read();p[i].v=v=read();
        p[i].lca=lca(u,v);
        p[i].dis=dis[u]+dis[v]-dis[p[i].lca]*2;
        cnt[u]++;
        v1[p[i].lca].push_back(u);
        v2[v].push_back(p[i].dis-dis[p[i].v]);
        v3[p[i].lca].push_back(p[i].dis-dis[p[i].v]);
    }
    dfs1(1);//自下至上
    dfs2(1);//自上至下
    for (i=1;i<=m;i++)
        if (dis[p[i].u]-dis[p[i].lca]==w[p[i].lca]) ans[p[i].lca]--;//去重
    for (i=1;i<=n;i++) printf("%d%c",ans[i],i<n?' ':'\n');
}