BOJ 1761 - 정점들의 거리

시간 제한	메모리 제한	제출	정답	맞은 사람	정답 비율
2 초	128 MB	6344	2550	1578	38.413%

문제

---

N(2 ≤ N ≤ 40,000)개의 정점으로 이루어진 트리가 주어지고 M(1 ≤ M ≤ 10,000)개의 두 노드 쌍을 입력받을 때 두 노드 사이의 거리를 출력하라.

입력

---

첫째 줄에 노드의 개수 N이 입력되고 다음 N-1개의 줄에 트리 상에 연결된 두 점과 거리를 입력받는다. 그 다음 줄에 M이 주어지고, 다음 M개의 줄에 거리를 알고 싶은 노드 쌍이 한 줄에 한 쌍씩 입력된다. 두 점 사이의 거리는 10,000보다 작거나 같은 자연수이다.

정점은 1번부터 N번까지 번호가 매겨져 있다.

출력

---

M개의 줄에 차례대로 입력받은 두 노드 사이의 거리를 출력한다.

접근

---

BOJ 11438 - LCA2
위 링크 문제를 응용한 문제였다.

M이 10,000이므로, 브루트 포스로 거리를 구하면 당연히 TLE가 난다.
그럼 어떻게 하면 하나의 쿼리를 계산하는 시간을 줄일 수 있을까?
답은 log 연산이다. 하나의 쿼리를 O(log(N))에 해결한다면, 모든 쿼리를 O(M * log(N))의 시간 복잡도로 해결할 수 있을 것이다.
루트 노드부터 DFS를 통해 각 노드의 깊이를 파악하고, 각 노드의 2^x 번째 조상 노드와 거기까지의 거리를 미리 계산해놓으면, log(N)의 시간에 쿼리를 해결할 수 있다.

이런 유형의 문제를 처음 접했다면, 시간이 많이 걸리지 않았을까 생각한다.

풀이

---

pair<int, int> par 2차원 배열은 [노드][지수] = { 조상 노드, 거리 }를 의미한다.
처음에 트리를 입력받으면, 1을 루트 노드로 잡고 DFS로 각 노드들의 깊이를 계산한다.
이 때, 2^0 번째, 즉 바로 부모 노드값도 함께 저장한다.
이 후 setPar 함수를 통해 지수를 점차 증가시키면서 조상 노드와 거리를 저장한다.

사전 작업이 끝나면, 쿼리를 입력받는다.
쿼리를 입력 받을때 마다, LCA 함수를 통해 거리를 측정한다.
우선 두 노드 x, y의 깊이를 맞춘다. 이 때, 해당 깊이의 노드가 같다면 바로 거리를 반환한다.
그 이후 log 연산을 통해 최소 공통 조상의 바로 아랫 노드까지 접근한다.
마지막으로, 접근한 노드에서 부모 노드까지의 거리를 각각 더해주면 정답이 반환된다.

#include <bits/stdc++.h>
using namespace std;
#define ll long long int
#define FUP(i, a, b) for(int i = a; i <= b; i++)
#define FDOWN(i, a, b) for(int i = a; i >= b; i--)
#define MS(a, b) memset(a, b, sizeof(a))
#define ALL(v) v.begin(), v.end()
#define CIN(a) cin >> a;
#define CIN2(a, b) cin >> a >> b
#define CIN3(a, b, c) cin >> a >> b >> c
#define COUT(a) cout << a
#define COUT2(a, b) cout << a << ' ' << b
#define COUT3(a, b, c) cout << a << ' ' << b << ' ' << c
#define ENDL cout << '\n'
int dy[4] = { -1, 1, 0, 0 };
int dx[4] = { 0, 0, 1, -1 };
#define MAX 100001
#define POW 17 //N의 최대 지수

int a, b, c, N, M, depth[MAX];
pair<int, int> par[MAX][POW + 1];
vector<pair<int, int>> edge[MAX];

void input()
{
	CIN(N);
	FUP(i, 1, N - 1)
	{
		CIN3(a, b, c);
		edge[a].push_back({ b, c });
		edge[b].push_back({ a, c });
	}
}

void DFS(int node, int before, int dist)
{
	par[node][0].first = before;
	par[node][0].second = dist;
	depth[node] = depth[before] + 1;
	for (auto next : edge[node])
	{
		if (next.first == before) continue;
		DFS(next.first, node, next.second);
	}
}

void setPar()
{
	FUP(j, 1, POW)
	{
		FUP(i, 1, N)
		{
			par[i][j].first = par[par[i][j - 1].first][j - 1].first;
			par[i][j].second = par[par[i][j - 1].first][j - 1].second + par[i][j - 1].second;
		}
	}
}

int LCA(int x, int y)
{
	int cnt = 0;
	if (depth[x] > depth[y]) swap(x, y);
	//x, y 깊이 맞추기
	FDOWN(i, POW, 0)
	{
		if (depth[y] - depth[x] >= (1 << i))
		{
			cnt += par[y][i].second;
			y = par[y][i].first;
		}
	}
	//일치할 경우 반환
	if (x == y) return cnt;
	//최소 공통 조상의 바로 아랫 노드로
	FDOWN(i, POW, 0)
	{
		if (par[x][i].first != par[y][i].first)
		{
			cnt += par[x][i].second + par[y][i].second;
			x = par[x][i].first;
			y = par[y][i].first;
		}
	}
	//부모 반환
	return cnt + par[x][0].second + par[y][0].second;
}

int main()
{
	ios_base::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);

	input();
	depth[0] = 0;
	DFS(1, 0, 0);
	setPar();

	CIN(M);
	while (M--)
	{
		CIN2(a, b);
		COUT(LCA(a, b));
		ENDL;
	}

	return 0;
}