[CodeUp 4035] 커피전문점

커피전문점

 

코드업을 밀던 중에 흥미로운 문제를 발견해서 포스팅한다.

다른 분들과 풀이의 방향이 달라 시간은 조금 느리지만, 다행히 코드는 깔끔하다. ㅎ

 

정해는 Union-Find 와 bfs 를 요구하지만, 사실 lca 만으로 풀린다!

생각의 편의를 위해 트리의 루트를 A로 만들어주자.

 

최적지의 집합이 (A-B 경로) $\cup $ (A-C 경로) 라는 것은 어렵지 않게 짐작 가능하다.

따라서 최적지의 개수는 $dep[a]+dep[b]-dep[lca(a,b)]+1$ 로 계산할 수 있다.

 

쿼리를 처리하는 것도 크게 복잡하지 않다.

후보지와 가장 가까운 최적지는 후보지에서 루트를 향해 올라가다가 처음으로 만나는 최적지이다.

이는 잘 생각해보면 B와의 LCA거나 A와의 LCA이고, 둘 중 루트에서 더 먼, 즉 깊이가 더 깊은 점이 정답임을 알 수 있다.

 

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+1;
vector<int> ad[N];
int dep[N],par[N][17];
void make_tree(int cur,int prv) {
	for(auto nxt : ad[cur]){
		if(nxt==prv) continue;
		dep[nxt]=dep[cur]+1, par[nxt][0]=cur;
		make_tree(nxt,cur);
	}
}
int lca(int u,int v) {
	if(dep[u]<dep[v]) swap(u,v);
	int diff=dep[u]-dep[v];
	for(int i=0;i<17;i++) if(diff&(1<<i)) u=par[u][i];
	if(u==v) return u;
	for(int i=16;i>=0;i--)
		if(par[u][i]!=par[v][i]) u=par[u][i], v=par[v][i];
	return par[u][0];
}
int main()
{
	ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
	int n; cin >> n;
	for(int i=1;i<n;i++){
		int u,v; cin >> u >> v;
		ad[u].push_back(v);
		ad[v].push_back(u);
	}
	int a,b,c; cin >> a >> b >> c;
	make_tree(a,0);
	for(int j=1;j<17;j++) for(int i=1;i<=n;i++)
		par[i][j]=par[par[i][j-1]][j-1];
	cout << dep[b]+dep[c]-dep[lca(b,c)]+1 << '\n';
	int q; cin >> q;
	while(q--){
		int v; cin >> v;
		int x=lca(v,b), y=lca(v,c);
		if(dep[x]<dep[y]) swap(x,y);
		cout << x << '\n';
	}
}