Codeforce 337D 容斥+树形dp

题目链接：http://codeforces.com/problemset/problem/337/D

题意：

给定n个点的树 m个点中有小怪物 距离l

m个点 表示该点有小怪物

n-1表示树

有一个u点存在大怪物，会把距离u点 <= l 的部分v点变成小怪物

告诉那些v点， 问所有可能的u点个数

即：问树中 距离M点集 <= l 的点个数

 

 

#include<stdio.h>#include<string.h>#include<vector>#include<iostream>using namespace std;#define N 100100vector<int>G[N];int dis[N];//dis[u] 表示以u为子树 距离u点最远的怪物距离 -1表示子树下无怪物int dp[N];//dp[u] 表示点集{所有点-以u为根的树} 中 距离最远的怪物 （无怪物=-1）   bool bird[N];//表示该点是否为小怪物void dfs1(int u, int fa){//计算 dis数组dis[u] = -1;for(int i = 0; i < G[u].size(); i++){int v = G[u][i]; if(v == fa)continue;dfs1(v, u);if(dis[v] != -1) dis[u] = max(dis[u], dis[v]+1);}if(bird[u])dis[u] = max(dis[u], 0);}void dfs2(int u, int fa){ //计算 dp数组int m1 = -1, m2 = -1;  //以u为根的子树 距离u点怪物最远和次远距离for(int i = 0; i < G[u].size(); i++){int v = G[u][i];if(v == fa || dis[v] == -1)continue;if(dis[v]+1 > m1)m2 = m1, m1 = dis[v] +1;else if(dis[v]+1 > m2)m2 = dis[v] +1;}for(int i = 0; i < G[u].size(); i++)//对于u的一个儿子v 实际上计算的是 u的其他儿子和v的关系{int v = G[u][i];    if(v == fa)continue;if( dis[v] +1 == m1 && dis[v] != -1)dp[v] = m2 ;else dp[v] = m1 ;////////////////这段计算 u 和 v 的关系if(dp[u] != -1) dp[v] = max(dp[v], dp[u]);if(bird[u]) dp[v] = max(dp[v], 0); /////////////if(dp[v] != -1)dp[v]++; //因为m1 m2的意义是 距离u点最远，而距离v点最远则是要+1dfs2(v, u);}}int main(){int u, v, i;int n, m, l;while(~scanf("%d %d %d",&n,&m,&l)){memset(bird, false, sizeof(bird));for(i = 1; i <= n; i++)G[i].clear();while(m--)scanf("%d",&u), bird[u] = true;for(i = 1; i < n; i++){scanf("%d %d",&u,&v);G[u].push_back(v);G[v].push_back(u);}dfs1(1, -1);dp[1] = -1;dfs2(1, -1);//设u点为大怪物可能点 则必有:所有小怪物距离u点 <=l    //容斥： [1]所有 以u为根的子树中的怪物 距离u点 <=l   //       [2]所有 不在u的子树中的点 中的怪物距离u点 <=l int ans = 0;for(i = 1; i <= n; i++)if(dis[i] <= l && dp[i] <= l)ans++;printf("%d\n",ans);}return 0;}