Codeforces 338B - Book of Evil

問題文 : Book of Evil

問題概要 :
大きさ$N$ の木で, $M $個の印をつけた頂点すべてからの距離が$d$ 以内であるものの個数を求めよ.

$N , M \leqq 10^5$

解法 :

印の付いた頂点のうち, 最遠点対を求める.
それらからの距離が両方$d$ 以内であれば, それはすべての印がついた頂点から$d$ 以内である.
(あとで証明を載せますが, 木の直径がこれで求まる原理とほぼ同じように示せます.)

コード :

#include <cstdio>
#include <vector>
#include <cstring>

using namespace std;

vector<int> G[100000];
int distu[100000], distv[100000];
bool marked[100000], vis[100000];

void dfs(int v, int *a, int d)
{
	vis[v] = true;
	a[v] = d;
	for (int i = 0; i < G[v].size(); i++)
		if (!vis[G[v][i]]) dfs(G[v][i], a, d + 1);
}

int main()
{
	int n, m, d;
	int t;
	scanf("%d %d %d", &n, &m, &d);
	
	for (int i = 0; i < m; i++){
		int x; scanf("%d", &x);
		marked[--x] = true;
		t = x;
	}
	
	for (int i = 0; i < n - 1; i++){
		int a, b;
		scanf("%d %d", &a, &b);
		G[--a].push_back(--b);
		G[b].push_back(a);
	}
	
	int x, y;
	dfs(t, distu, 0);
	int dd = -1;
	for (int i = 0; i < n; i++) if (marked[i] && distu[i] > dd){dd = distu[i]; x = i;}
	memset(vis, 0, sizeof(vis));
	dfs(x, distu, 0);
	dd = -1;
	memset(vis, 0, sizeof(vis));
	for (int i = 0; i < n; i++) if (marked[i] && distu[i] > dd){dd = distu[i]; y = i;}
	dfs(y, distv, 0);
	
	int ans = 0;
	for (int i = 0; i < n; i++) if (distu[i] <= d && distv[i] <= d) ans++;
	
	printf("%d\n", ans);
	
	return (0);
}