题目大意:树上某点到所有叶超能力点的距离和超能力节点的距离的和
思路:1.先用bfs求出各点到最近的超能力节点的距离并标记叶超能力节点
            2.用dfs1从根节点开始从上向下遍历求出到每个节点到叶超能力节点的距离和(sum0z数组)、到所有超能力节点的距离和(sum1数组)、到所有超能力节点距离的平方和(数组sum2)
            3.通过父节点到所有超能力节点距离的平方的状态转移到子节点到所有超能力点距离的平方的状态
void dfs2( int x,int fa )
{
    ans[x]=sum2[x]+sum0[x];
    for(int i=0;i<g[x].size();++i )
    {
        int v=g[x][i];
        if( v==fa )
            continue;
        ll tmp=sum1[x]-sum1[v]-cnt1[v];
        sum2[v]=sum2[x]-2*(sum1[v]+cnt1[v])+tmp*2+m;
        sum1[v]+=(tmp+n-cnt1[v]);
        sum0[v]+=(sum0[x]-sum0[v]-cnt2[v]+cnt2[1]-cnt2[v]);
        dfs2(v,x);
    }
}
这个函数有点难理解,详细讲解一下


以图中通过sum1[1]求sum1[2]为例进行讲解
/*dis数组为父节点1到各个超能力节点的距离
sum1[1]=dis[3]*dis[3]+dis[9]*dis[9]
sum1[2]=(dis[3]+1)*(dis[3]+1)+(dis[9]-1)*(dis[9]-1)=dis[3]*dis[3]+dis[9]*dis[9]+2*dis[3]-2*dis[9]+1+1;
            =sum1[1]+2*dis[3]-2*dis[9]+1+1;
现在推广到一般:dis[3]要理解成父节点到右子树所有超能力节点距离的平方,dis[9]要推广到父节点1到
子节点2其他所有超能力点的距离和,最后面的两个1分别表示父节点左子树所有的超能力的点和父节点右子树所有超能力的点的数量,所以加起来就是超能力点的总数
dfs1 的代码中tmp表示的就是dis[1],sum1[v]+cnt1[v]表示的就是dis[9](注意sum1[v]是下一行才更新,所以表示的是2节点到下方超能力结点距离的平方)
叶能力节点的距离更新=父节点到左子树所有叶超能力节点的距离和+子节点到下方叶能力节点的距离和+
子节点下方叶能力节点的个数(父节点要经过多少次起始点为父节点和子节点的边)
其实这个题目跟牛客上的 牛客练习赛55-树 的升级版
上题详解点击下方大佬博客
AC代码
#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const int maxn=5e5+10;
vector<int> g[maxn];
int sup[maxn],psup[maxn],vis[maxn];
int n,m,k,l,r;
int dis[maxn];
bool vsup[maxn],vpsup[maxn];
ll sum1[maxn],sum2[maxn]; // 超能力距离之和  超能力距离平方和
ll sum0[maxn]; // 旪超能力距离之和
ll cnt1[maxn],cnt2[maxn];  // 从第i个节点从上向下遍历得到的超能力节点和叶节点能力(包括自身)
ll ans[maxn];
void bfs()
{
    queue<int> que;
    for( int i=1; i<=m; i++ )
        vis[sup[i]]=1,vsup[sup[i]]=true,dis[sup[i]]=0,que.push(sup[i]);
    while( !que.empty() )
    {
        int p=que.front();
        que.pop();
        for( int i=0;i<g[p].size();++i )
        {
            int v=g[p][i];
            if( vis[v] )
                continue;
            dis[v]=dis[p]+1;///bfs是按距离从小到大扫的
            vis[v]=1;
            que.push(v);
        }
    }
    for( int i=1; i<=n; i++ )
        if( dis[i]>=l && dis[i]<=r )//第i个点到超能力点的最短距离
            vpsup[i]=true;
}
void dfs1( int x,int fa )
{
    for(int i=0;i<g[x].size();++i)
    {
        int v=g[x][i];
        if( v==fa )
            continue;
        dfs1(v,x);
        cnt1[x]+=cnt1[v];
        cnt2[x]+=cnt2[v];
        sum1[x]+=sum1[v]+cnt1[v];
        sum2[x]+=sum1[v]*2+sum2[v]+cnt1[v];
        sum0[x]+=sum0[v]+cnt2[v];
    }               ///dfs2前的sum0,sum1,sum2代表从第i个节点从上往下开始遍历
    if( vsup[x] )   ///到每个叶超能力节点的距离和到每个超能力节点距离和距离平方和
        cnt1[x]++;
    if( vpsup[x] )
        cnt2[x]++;
}
void dfs2( int x,int fa )
{
    ans[x]=sum2[x]+sum0[x];
    for(int i=0;i<g[x].size();++i )
    {
        int v=g[x][i];
        if( v==fa )
            continue;
        ll tmp=sum1[x]-sum1[v]-cnt1[v];
        sum2[v]=sum2[x]-2*(sum1[v]+cnt1[v])+tmp*2+m;
        sum1[v]+=(tmp+cnt1[1]-cnt1[v]);
        sum0[v]+=(sum0[x]-sum0[v]-cnt2[v]+cnt2[1]-cnt2[v]);
        dfs2(v,x);
    }
}
int main()
{
    scanf("%d%d%d%d%d",&n,&m,&k,&l,&r);
    for( int i=1; i<n; i++ )
    {
        int x,y;
        scanf("%d%d",&x,&y);
        g[x].push_back(y);
        g[y].push_back(x);
    }
    for( int i=1; i<=m; i++ )
        scanf("%d",&sup[i]);
    sort(sup+1,sup+1+m);
    m=unique(sup+1,sup+1+m)-(sup+1);
    bfs();
    dfs1(1,0);
    dfs2(1,0);
    while( k-- )
    {
        int x;
        scanf("%d",&x);
        printf("%lld\n",ans[x]);
    }
}
注:此代码借鉴了牛友
一命为红颜
的AC代码