2020牛客NOIP赛前集训营-普及组第三场C牛半仙的妹子树

链接：https://ac.nowcoder.com/acm/contest/7608/C
来源：牛客网
 

牛半仙有 n​ 个妹子，她们所在的位置组成一棵树，相邻两个妹子的距离为 1​。
有 m​ 个妹子具有超能力，她们会影响到其他妹子。
离所有具有超能力的妹子的最短距离在 [l,r]间的妹子会受到影响，会具有旪超能力。
这些具有能力的妹子共同形成了一个磁场。对于一个位置，一个具有超能力的妹子为其增加的磁场强度为妹子到这个位置的距离的平方，一个具有旪超能力的妹子为其增加的磁场强度为妹子到这个位置的距离。
现在牛半仙想知道一个位置的磁场强度有多大。

因为牛半仙对妹子们特别关心，所以他有 k​ 个询问。

输入描述:

第一行 5个正整数，代表 n,m,k,l,r​ 。
第 2​ 行到第 n行每行 2​ 个正整数 ui​,vi​，代表树中存在边 (ui,vi)。
第 n+1​ 行有 m​ 个正整数，代表这些妹子有超能力。
第 n+2​ 行有 k​ 个正整数，代表询问的妹子。对于每个询问，输出一行，代表该询问的答案。

输出描述:

对于每个询问，输出一行，代表该询问的答案。

示例1

输入

11 2 4 2 2
1 2
1 3
2 4
2 5
3 6
3 7
5 8
5 9
9 10
9 11
3 9
1 6 9 11

输出

14
34
20
32

题目大意：树上某点到所有叶超能力点的距离和超能力节点的距离的和

思路：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-树 的升级版

上题详解点击下方大佬博客

https://blog.csdn.net/qq_21433411/article/details/103534413

#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代码