Codeforces #196 Div1. B. Book of Evil

Bでグラフ問題か…。
http://codeforces.com/contest/338/problem/B

問題

木を成したN点からなる無向グラフが与えられる。
N点のうち、M点が指定されるので、M点すべてからの距離がD以内の点の数を答えよ。

解法

M点のうち１つをrootとして、２回DFSしていけばよい。
１回目で、各点から伸びる各辺において、その先にM点に選ばれた点があればそこまでの距離の最大値を求める。
２回目のDFSでは、辺をたどる際、親からの距離と他の辺からの距離の最大値を子に伝えていきながら、各点がM点から距離D以内か判定すればよい。

int N,M,D;
int P[1000001],R;
vector<pair<int,int> > E[1000001];

int dfs(int cur,int pre) {
	int re=-1000000,i,x,p=-1;
	if(P[cur]==D+1) re=0;
	
	FOR(i,E[cur].size()) {
		int tar=E[cur][i].second;
		if(tar==pre){ p=i; continue;}
		E[cur][i].first = dfs(tar,cur)+1;
		re=max(re,E[cur][i].first);
	}
	if(p>=0) E[cur].erase(E[cur].begin()+p);
	return re;
}

void dfs2(int cur,int de2) {
	int i,ng=0;
	if(de2<0) return;
	
	sort(E[cur].begin(),E[cur].end());
	FOR(i,E[cur].size()) if(E[cur][i].first>D && E[cur][i].first!=10000000) ng++;
	if(ng==0) R++;
	
	if(E[cur].size()==1) dfs2(E[cur][0].second,de2-1);
	else if(E[cur].size()>1) {
		FOR(i,E[cur].size()-1) dfs2(E[cur][i].second,min(de2-1,D-E[cur][E[cur].size()-1].first-1));
		dfs2(E[cur][E[cur].size()-1].second,min(de2-1,D-E[cur][E[cur].size()-2].first-1));
	}
}

void solve() {
	int f,i,j,k,l, x,y;
	ll ret=0;
	
	cin>>N>>M>>D;
	FOR(i,M) P[GETi()-1]=D+1;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(make_pair(0,y-1));
		E[y-1].push_back(make_pair(0,x-1));
	}
	
	FOR(i,N) if(P[i]==D+1) break;
	x=i;
	dfs(x,-1);
	FOR(i,N) FOR(j,E[i].size()) if(E[i][j].first<0) E[i][j].first=0;
	dfs2(x,D);
	_P("%d\n",R);
}

まとめ

Codeforcesはグラフを複数回たどることでO(N)で処理する、って問題が多いな。
このスタイルにだいぶ慣れてきた気がする。