취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 25198번 문제인 곰곰이의 심부름이다.
문제는 아래 링크를 확인하자.

https://www.acmicpc.net/problem/25198 

중간지점인 C를 기준으로, C에서 S, C에서 H로 가는 경로들을 각각 LS, LH라 하자.

LS와 LH의 공통노드 중 C에서 가장 먼 노드를 X라 할 때, 문제의 경로는 "S에서 X이전까지", "X에서 C를 찍고 다시 X까지", "X 다음서부터 H까지"의 세 부분으로 나눌 수 있다. 이러한 X는 C를 루트로 하는 트리에서의 S와 H의 LCA로 구해낼 수 있다.

각 부분의 노드 개수를 이용하여 답을 계산해내자.

아래는 제출한 소스코드이다.

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
typedef long long ll;

vector<int> G[100001];
int level[100001];
int par[100001];
void dfs(int cur) {
	for (auto& node : G[cur]) {
		if (level[node]) continue;
		level[node] = level[cur] + 1;
		par[node] = cur;
		dfs(node);
	}
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	int N; cin >> N;
	int S, C, H; cin >> S >> C >> H;
	for (int i = 1; i < N; i++) {
		int u, v; cin >> u >> v;
		G[u].emplace_back(v);
		G[v].emplace_back(u);
	}

	level[C] = 1;
	dfs(C);
	if (level[S] < level[H]) swap(S, H);
	ll cntS = level[S], cntH = level[H];

	ll tmpS = cntS, tmpH = cntH;
	while (tmpS > tmpH) {
		S = par[S];
		tmpS--;
	}

	while (S != H) {
		tmpS--; tmpH--;
		S = par[S];
		H = par[H];
	}

	ll cntLCA = level[S];
	cntS -= cntLCA, cntH -= cntLCA;
	cout << cntLCA * (cntLCA - 1) + (cntS + cntH) * cntLCA + cntS * cntH + (cntS * (cntS - 1)) / 2 + (cntH * (cntH - 1)) / 2;
}