1. 程式人生 > >337D Book of Evil (樹形dp)

337D Book of Evil (樹形dp)

題意:有一顆有n個結點的樹,樹上存在一個汙染源(位置不確定),它可以汙染與它距離不超過d的節點,現給出m個被汙染的節點(汙染源本身也可能是被汙染的節點),求汙染源可能的位置數。

思路:樹形dp,樹的最長鏈,dp[i][0]表示i到i的子樹中最遠汙染源的距離,dp[i][1]表示i到i的子樹中次遠汙染源的距離,dp[i][2]表示i經過i的父節點到最遠汙染源的距離。那麼汙染源到這些被汙染的點的最遠距離應該不大於d,列舉一遍統計合法的點即為答案。

其他細節看程式碼。

#include<bits/stdc++.h>
using namespace std;
int dp[100005][3];//dp[i][0]表示i到i的子樹中最遠汙染源的距離,dp[i][1]表示i到i的子樹中次遠汙染源的距離,dp[i][2]表示i經過i的父節點到最遠汙染源的距離.
bool vis[100005];
vector<int>vec[100005];
void dfs1(int x,int y) {
    int k=vec[x].size();
    for(int i=0; i<k; i++) {
        int v=vec[x][i];
        if(v==y)
            continue;
        if(vis[v])//該點被汙染更新dp[v][0]
            dp[v][0]=0;
        dfs1(v,x);//先更新子樹中的答案
        //dp[v][0]>=0說明v的子樹中存在被汙染的點,更新父結點的答案
        if(dp[v][0]>=0) {
            if(dp[v][0]+1>=dp[x][0]) {
                dp[x][1]=dp[x][0];
                dp[x][0]=dp[v][0]+1;
            } else if(dp[v][0]+1>dp[x][1])
                dp[x][1]=dp[v][0]+1;
        }
    }
}
void dfs2(int x,int y) {
    int k=vec[x].size();
    //x為被汙染點,應更新答案,但存在x的父結點中有被汙染點已經更新答案,所以取最大的答案
    if(vis[x])
        dp[x][2]=max(0,dp[x][2]);
    for(int i=0; i<k; i++) {
        int v=vec[x][i];
        if(v==y)
            continue;
        int ans=max(dp[x][2],(dp[x][0]==dp[v][0]+1)?dp[x][1]:dp[x][0]);
        //答案合法,更新子節點的答案
        if(ans>=0)
            dp[v][2]=ans+1;
        dfs2(v,x);
    }
}
int main() {
    int n,m,d;
    scanf("%d%d%d",&n,&m,&d);
    for(int i=0; i<m; i++) {
        int a;
        scanf("%d",&a);
        vis[a]=true;
    }
    for(int i=0; i<n-1; i++) {
        int a,b;
        scanf("%d%d",&a,&b);
        vec[a].push_back(b);
        vec[b].push_back(a);
    }
    for(int i=1; i<=n; i++) {
        dp[i][0]=dp[i][1]=dp[i][2]=-10;
    }
    dfs1(1,0);
    dfs2(1,0);
    int sum=0;
    //cout<<dp[2][2]<<endl;
    for(int i=1; i<=n; i++) {
        if(max(dp[i][0],dp[i][2])<=d)
            sum++;
    }
    printf("%d\n",sum);
    return 0;
}