yukicoder : No.1718 Random Squirrel

難しくはないけど、割と面倒。
https://yukicoder.me/problems/no/1718

問題

木を成すグラフが与えられる。うち、いくつかの頂点にはドングリが置いてある。
ある点から辺をたどって移動し、ドングリがある頂点をそれぞれ最低1回通るようにしたい。

各点を始点としたとき、条件を満たす最小移動回数はいくつか。

解法

ドングリが置いてある点からなる最小シュタイナー木を考える。
もし、ドングリのある点を全部回収して始点に戻る場合、どこが始点でも、その移動回数は(辺の数)×2で一致する。
ただ、この問題では実際には始点に戻らなくても良い。
ということは、最後のドングリのある点から始点に戻る移動を省略できると考えると、最遠点のドングリを最後に取るのが最適である。

そこで、シュタイナー木内の各点から、最遠点までの距離を求めよう。
これでシュタイナー木内の頂点に関する解が得られる。
シュタイナー木外の頂点については、まず最寄のシュタイナー木内の頂点に移動することを考えると解は自明。

int N,K;
vector<int> E[202020];
vector<int> D[2];
int G[202020];
int dep[202020],from[202020];
int ret[202020];

int dfs(int cur,int pre) {
	FORR(e,E[cur]) if(e!=pre) G[cur]|=dfs(e,cur);
	return G[cur];
}

pair<int,int> farthest(int cur,int pre,int d,vector<int>& D) {
	D[cur]=d;
	pair<int,int> r={d,cur};
	FORR(e,E[cur]) if(e!=pre&&G[e]) r=max(r, farthest(e,cur,d+1,D));
	return r;
}

//両端と間の頂点を返す
vector<pair<int,int>> diameter(int root) { // diameter,center
	D[0].resize(N);
	D[1].resize(N);
	auto v1=farthest(root,root,0,D[0]);
	auto v2=farthest(v1.second,v1.second,0,D[0]);
	farthest(v2.second,v2.second,0,D[1]);
	
	vector<pair<int,int>> R;
	for(int i=N-1;i>=0;i--) if(D[0][i]+D[1][i]==v2.first) R.push_back({D[0][i],i});
	sort(ALL(R));

	return R;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N-1) {
		cin>>x>>y;
		E[x-1].push_back(y-1);
		E[y-1].push_back(x-1);
	}
	FOR(i,K) {
		cin>>x;
		G[x-1]=1;
	}
	int root;
	FOR(i,N) if(G[i]==1) root=i;
	dfs(root,root);
	
	queue<int> Q;
	FOR(i,N) {
		if(G[i]==1) from[i]=i, Q.push(i);
		else dep[i]=1<<20;
	}
	while(Q.size()) {
		int cur=Q.front();
		Q.pop();
		FORR(e,E[cur]) if(dep[e]>dep[cur]+1) {
			dep[e]=dep[cur]+1;
			from[e]=from[cur];
			Q.push(e);
		}
	}
	
	auto v=diameter(root);
	int len=0;
	FOR(i,N) if(G[i]) FORR(e,E[i]) len+=G[e];
	FOR(i,N) if(G[i]) {
		ret[i]=len-max(D[0][i],D[1][i]);
	}
	FOR(i,N) {
		if(G[i]==0) ret[i]=dep[i]+ret[from[i]];
		cout<<ret[i]<<endl;
	}
	
	
		
}

まとめ

典型テクの詰め合わせ感のある問題。