CodeForces - 337D D. Book of Evil

题面

题意

给出一棵树,其中有且只有一个节点上有魔鬼,与它距离在d以内的点都有可能收到影响,先给你这些收到影响的点,问几个节点上可能有魔鬼

做法

首先可以证明:如果有一点到这些点中的最远的两点的距离均小于等于d,则这点上可能会有鬼.
之后问题就转化为了找收到影响的点中的最远两点.
又可以证明:以任意一点为根,两点中一定有一个点的深度最大.
因此可以先以1为根,dfs,找到此时深度最大的且受到影响点,以它为根,dfs,找到另外一个点,再以另一个点为根,dfs,那么满足二三两次dfs的深度均小于等于d的点均可能有鬼.

代码

#include<iostream>
#include<cstdio>
#include<cstring>
#define N 100100
using namespace std;

int n,m,d,deep[4][N],num[N],bb,first[N],ans;
struct Bn
{
    int to,next;
}bn[N*2];

inline void add(int u,int v)
{
    bb++;
    bn[bb].to=v;
    bn[bb].next=first[u];
    first[u]=bb;
}

void dfs(int u,int now,int last)
{
    int p,q;
    for(p=first[now];p!=-1;p=bn[p].next)
    {
        if(bn[p].to==last) continue;
        deep[u][bn[p].to]=deep[u][now]+1;
        dfs(u,bn[p].to,now);
    }
}

int main()
{
    memset(first,-1,sizeof(first));
    int i,j,p,q,mm,mx;
    cin>>n>>m>>d;
    for(i=1;i<=m;i++)
    {
        scanf("%d",&num[i]);
    }
    for(i=1;i<n;i++)
    {
        scanf("%d%d",&p,&q);
        add(p,q);
        add(q,p);
    }
    dfs(1,1,-1);
    mx=0;
    for(i=1;i<=m;i++)
    {
        if(deep[1][num[i]]>mx)
        {
            mx=deep[1][num[i]];
            mm=num[i];
        }
    }
    dfs(2,mm,-1);
    mx=0;
    for(i=1;i<=m;i++)
    {
        if(deep[2][num[i]]>mx)
        {
            mx=deep[2][num[i]];
            mm=num[i];
        }
    }
    dfs(3,mm,-1);
    for(i=1;i<=n;i++)
    {
        if(deep[2][i]<=d&&deep[3][i]<=d)
        {
            ans++;
        }
    }
    cout<<ans;
}