Codeforces 338B Book of Evil

問題

n頂点からなる木のうち、m個の頂点から被害が出ている。
頂点のうち1箇所に悪霊がいて、被害の出ている頂点は全て、悪霊のいる頂点からの距離がd以内である。

このとき、悪霊のいる頂点としてありうる頂点の個数を求めよ。

制約条件

n, m, d≦10^5

方針

悪霊のいる頂点の最遠点対(u, v)を一つ求める。
これは木の直径を求めるような感じで、DFSorBFSを二回やればよい。

そしたら、悪霊がいる頂点というのは、uからvへのパス上の頂点からちょっとずれたとこが候補になりそうである。
で、どれだけずれて良いかというと、d - max(dist(cur, u), dist(cur, v))
つまり、遠いほうからx遠ざかってもdの距離に入ってるなら必要かつ十分。
なので、パス上の頂点をキューに入れて幅優先探索すればいい。

ソースコード

int n, m, d, prev[100000];
bool mark[100000], v[100000];
vector<vi> e;

pi calc(int s){
	pi res(-1, -1);
	queue<pi> q;
	
	memset(prev, -1, sizeof(prev));
	prev[s] = -2;
	q.push(mp(s, 0));
	while(!q.empty()){
		int c = q.front().first, co = q.front().second;
		q.pop();
		if(mark[c]) res = max(res, mp(co, c));
		
		rep(i, e[c].size()) if(prev[e[c][i]] == -1){
			prev[e[c][i]] = c;
			q.push(mp(e[c][i], co + 1));
		}
	}
	return res;
}

int main(){
	cin >> n >> m >> d;
	e.resize(n);
	
	int a;
	rep(i, m){
		cin >> a;
		mark[--a] = 1;
		dbg(a);
	}
	rep(i, n - 1){
		int a, b;
		cin >> a >> b;
		a--; b--;
		e[a].pb(b);
		e[b].pb(a);
	}
	
	int ans = 0;
	queue<pi> q;
	pi p = calc(a);
	p = calc(p.second);
	
	for(int c = p.second, cnt = 0; c >= 0; c = prev[c], cnt++){
		q.push(mp(c, d - max(cnt, p.first - cnt)));
	}
	while(!q.empty()){
		int c = q.front().first, co = q.front().second;
		q.pop();
		if(co < 0 || v[c]) continue;
		v[c] = 1;
		
		ans++;
		if(co > 0) rep(i, e[c].size()) if(!v[e[c][i]]) q.push(mp(e[c][i], co - 1));
	}
	
	cout << ans << endl;
	
	return 0;
}