Codeforces 337D Book of Evil 树状DP 或 BFS找子树直径端点

题目大意：

就是现在对于一个n个点的树 (1 <= n <= 10^5), 其中有m个点出现了恶魔(1 <= m <= n), 给出m个点的编号(p1, p2, ... pm) 1 <= pi <= n, 现在已知这些恶魔出现的原因是有一本恶魔之书造成的, 而且书到所有恶魔出现的点的距离不能超过d (0 <= d <= n - 1), 问这棵树中有哪些点可能是恶魔之书出现的位置, 输出可能的位置的数量

大致思路：

这个题试了两种解法, 第一种是官方题解的树状DP, 另外一种是比较巧妙的转化成树的直径端点相关的问题

细节都写在代码注释里了, 详情见代码吧

代码如下：

解法一: 树状DP解法

Result  :  Accepted     Memory  :  7504 KB     Time  :  216 ms

/*
 * Author: Gatevin
 * Created Time:  2015/3/4 20:09:20
 * File Name: Kotori_Itsuka.cpp
 */
#include<iostream>
#include<sstream>
#include<fstream>
#include<vector>
#include<list>
#include<deque>
#include<queue>
#include<stack>
#include<map>
#include<set>
#include<bitset>
#include<algorithm>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cctype>
#include<cmath>
#include<ctime>
#include<iomanip>
using namespace std;
const double eps(1e-8);
typedef long long lint;

/*
 * 解法一: 树状DP
 * 首先任意选择一个点作为树的根节点, 方便起见我们选择点1作为根
 * 用集合P表示被标记的点的集合
 * 用disDown[i][0]表示从点i到位于以i为根的子树上的被标记的点的距离的最大值
 * 用disDown[i][1]表示从点i到位于i为根的子树上的被标记的点的距离的次大值
 * 注意这里如果disDown[i][0]与disDown[i][1]如果同时存在,
 * 那么产生这两个值得被标记的点一定在i的两颗不同的子树上
 * 用disUp[i]表示点i到不位于以i为根的子树上的P中的点的距离的最大值
 * 注意disUp[i]的值如果存在其最短路径一定经过i的父亲节点
 * 那么满足题意的点药满足的条件就是disDown[i][0] <= d && disUp[i] <= d
 * disDown[i][0]与disDown[i][1]可以一遍dfs弄出来
 * disUp满足当u是v的父亲节点时有
 * 如果disDown[u][0] = disDown[v][0] + 1说明造成disDown[u][0]的点位于v所在子树上
 * 此时, disDown[u][1]的来源点一定是v的兄弟所在树上的点造成的(如果存在的话)
 * 那么disUp[v] = max(disUp[u] + 1, disDown[u][1] + 1)
 * 否则的话说明disDown[u][0]来自v以外的兄弟所在子树, 由于disDown[u][0] >= disDown[u][1]
 * 有disUp[v] = max(disUp[u] + 1, disDown[u][0] + 1)
 * 于是disUp也可以一遍dfs解决
 * 总体复杂度O(n)
 */
const int inf = 0x3f3f3f3f;
int n, m, d;
bool evil[100010];
vector <int> G[100010];
int disDown[100010][2], disUp[100010];

void dfs_disDown(int now, int father)
{
    if(evil[now]) disDown[now][0] = disUp[now] = 0;
    for(unsigned int i = 0, sz = G[now].size(); i < sz; i++)
    {
        int nex = G[now][i];
        if(nex == father) continue;
        dfs_disDown(nex, now);
        if(disDown[now][0] < disDown[nex][0] + 1)//记录第一第二大值
        {
            disDown[now][1] = disDown[now][0];
            disDown[now][0] = disDown[nex][0] + 1;
        }
        else disDown[now][1] = max(disDown[now][1], disDown[nex][0] + 1);
    }
    return;
}

void dfs_disUp(int now, int father)
{
    for(unsigned int i = 0, sz = G[now].size(); i < sz; i++)
    {
        int nex = G[now][i];
        if(nex == father) continue;
        if(disDown[now][0] == disDown[nex][0] + 1)//造成dis[now][0]的来自nex所在子树
            disUp[nex] = max(disUp[now] + 1, disDown[now][1] + 1);//比较第二大值
        else disUp[nex] = max(disUp[now] + 1, disDown[now][0] + 1);//只需要比较第一大值
        dfs_disUp(nex, now);
    }
    return;
}

int main()
{
    scanf("%d %d %d", &n, &m, &d);
    for(int i = 1; i <= n; i++)
        disDown[i][0] = disDown[i][1] = -inf;//注意初始化
    fill(disUp, disUp + n + 1, -inf);
    int tmp;
    while(m--)
        scanf("%d", &tmp), evil[tmp] = 1;
    int tu, tv;
    for(int i = 1; i < n; i++)
    {
        scanf("%d %d", &tu, &tv);
        G[tu].push_back(tv);
        G[tv].push_back(tu);
    }
    dfs_disDown(1, 0);
    dfs_disUp(1, 0);
    int ans = 0;
    for(int i = 1; i <= n; i++)
        if(disDown[i][0] <= d && disUp[i] <= d)
            ans++;
    printf("%d\n", ans);
    return 0;
}

解法二:  转化为到树的直径端点距离的问题

Result  :  Accepted     Memory  :  5596 KB     Time  :  186 ms

/*
 * Author: Gatevin
 * Created Time:  2015/3/4 21:50:07
 * File Name: Kotori_Itsuka.cpp
 */
#include<iostream>
#include<sstream>
#include<fstream>
#include<vector>
#include<list>
#include<deque>
#include<queue>
#include<stack>
#include<map>
#include<set>
#include<bitset>
#include<algorithm>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cctype>
#include<cmath>
#include<ctime>
#include<iomanip>
using namespace std;
const double eps(1e-8);
typedef long long lint;

/*
 * 解法二: 求包含所有被标记的点的最小子树的直径
 * 然后比较其他点直径的两个端点的距离即可
 * 注意到有这样一个事实:
 * 如果一棵树T的直径上的两个端点分别是A, B
 * 且T是树S的一部分
 * 那么如果S上某个点到A, B的距离不超过D
 * 那么这个点到这棵子树上的所有点的距离不超过D
 * 所以只需要找出包含所有点P[1~m]的最小的树T之后
 * 判断其他点到这棵树直径上的两个端点的距离是否 <= d即可
 */

int n, m, d;
vector <int> G[100010];
int d1[100010], d2[100010];
int p[100010];

int bfs(int start)
{
    queue <int> Q;
    memset(d1, -1, sizeof(d1));
    d1[start] = 0;
    Q.push(start);
    while(!Q.empty())
    {
        int now = Q.front();
        Q.pop();
        for(unsigned int i = 0, sz = G[now].size(); i < sz; i++)
            if(d1[G[now][i]] == -1)
                d1[G[now][i]] = d1[now] + 1, Q.push(G[now][i]);
    }
    int ret = 1;
    for(int i = 1; i <= m; i++)
        if(d1[p[i]] > d1[p[ret]]) ret = i;
    return p[ret];
}

int main()
{
    scanf("%d %d %d", &n, &m, &d);
    for(int i = 1; i <= m; i++)
        scanf("%d", p + i);
    int u, v;
    for(int i = 1; i < n; i++)
    {
        scanf("%d %d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    int A = bfs(p[1]);//A是直径的一端
    int B = bfs(A);//B是直径的另外一端
    for(int i = 1; i <= n; i++)
        d2[i] = d1[i];
    //d2[i]为点i到A的距离
    bfs(B);
    //d1[i]为点i到B的距离
    int ans = 0;
    for(int i = 1; i <= n; i++)
        if(d1[i] <= d && d2[i] <= d)
            ans++;
    printf("%d\n", ans);
    return 0;
}