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 100100
vector<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点最远则是要+1
		dfs2(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;
}