boj, 백준) 13511. 트리와 쿼리2 ( C / C++)

1. 문제 링크

www.acmicpc.net/problem/13511

13511번: 트리와 쿼리 2

 

2. 문제 개요

N개의 정점으로 이루어진 트리가 있다. (트리 : connected acyclic undirected graph ) 정점은 1번부터 N번까지 번호가 매겨져 있다.

 

두 개의 쿼리를 수행하는 프로그램을 작성하자.

 

(1) 1 u v : u에서 v로가는 경로의 비용을 출력한다.

(2) 2 u v k : u에서 v로 가는 경로에 존재하는 정점 중 k번째 정점을 출력하자.

 

 

3. 문제 힌트

트리에서 경로는 LCA로 구할 수 있다. 물론 BFS/DFS 같은 걸로도 구할 수 있지만.. 너무 느리다.

정점이 10^5개, 쿼리가 10^5개이다.

BFS/DFS경우 쿼리당 O(n)만큼의 탐색을 할 수 있으므로, O(mn)으로는 불가.

 

LCA는 2^x만큼 뛰어가며 탐색을하기 때문에 log2(n)의 연산으로 거리를 찾을 수 있다. 트리에서 (최단) 거리는 LCA를 찍고 간다는 것을 활용.

 

4. 문제 풀이

다른분들의 풀이는 Root->u까지의 비용 + root->v까지의 비용 - 2*root->LCA까지 비용을 사용하여 구했던데 저는 parent배열과 같이 1, 2, 4, 8,... 부모까지 가는데 드는 비용을 저장해서 풀었습니다.

 

위의 root에서 정점 x까지 가는 비용은 DFS함수에서 누적하며 구하면 됩니다.

 

 

풀면서 위의 아이디어를 생각 못하긴 했지만, 어차피 parent배열을 채워나가면서 cost배열도 똑같이 채워나가면 된다고 생각했기 때문에 LCA 구하면서 함께 구할 수 있다고 생각했습니다.

 

그래서, LCA를 구할때 정점을 위로 올려가면서 비교할 텐데, 그때 cost값도 함께 더해주면서 LCA와 비용(u->LCA로 가는 비용 + v->LCA로 가는 비용)을 구했습니다.

 

(2)번 쿼리 같은 경우는 종이에 트리를 그려서 구하는 것이 생각하기에 엄청 좋습니다.

depth정보를 사용했습니다.

 

 

2번 쿼리

위 처럼 depth를 구해서 개수를 구합니다. 그리고 LCA를 기준으로 왼쪽 오른쪽을 살펴보고, k가 어느 subtree에 위치하는지 찾아냅니다.

 

왼쪽 같은 경우는 단순히 출발 정점에서 k-1칸(자기 자신 포함이기 때문)만큼 위로 올라가면 됩니다.

오른쪽이 헷갈리는데.. 출발정점 -> LCA -> 도착 정점에서 출발정점->LCA - 1만큼 제외합니다.

그러면 위 그림의 초록색 동그라미 부분만 남게 되는데,

예를들어 초록색 동그라미 부분에 10개의 정점이 있는데 여기서 3번째를 찾으려고 한다, 라고 하면 도착 정점에서 7칸을 올라가는 방식으로 구할 수 있다.

 

5. 코드

#include <cstdio>
#include <algorithm>
#include <vector>
#include <cmath>
using namespace std;

int n, m, len;
vector<vector<pair<int, int>>> adj;   //to, weight

//for LCA
int depth[100001], parent[100001][17];
long long cost[100001][17];

pair<int, long long > LCA(int u, int v) {

	long long ret = 0;

	if (depth[u] < depth[v])
		swap(u, v);

	int diff = depth[u] - depth[v];

	for (int k = 0; diff != 0; ++k) {

		if (diff % 2 == 1) {
			ret += cost[u][k];
			u = parent[u][k];
		}
		diff /= 2;
	}


	if (u != v) {

		for (int k = len - 1; k >= 0; --k) {
			if (parent[u][k] != parent[v][k]) {
				ret += cost[u][k] + cost[v][k];
				u = parent[u][k];
				v = parent[v][k];
			}
		}
		ret += cost[u][0] + cost[v][0];
		u = parent[u][0];
	}

	return { u,ret };
}

void dfs(int cur, int before)
{
	for (pair<int, long long> next : adj[cur])
	{
		if (next.first == before)
			continue;

		parent[next.first][0] = cur;
		cost[next.first][0] = next.second;
		depth[next.first] = depth[cur] + 1;
		dfs(next.first, cur);
	}
	return;
}


int main()
{
	scanf("%d", &n);
	adj.resize(n + 1);
	len = (int)ceil(log2(n));

	for (int i = 0; i < n - 1; ++i) {
		int from, to, weight;

		scanf("%d %d %d", &from, &to, &weight);

		adj[from].push_back({ to,weight });
		adj[to].push_back({ from,weight });
	}

	dfs(1, 0);   //root is a 'vertex 1'

	for (int k = 0; k < len - 1; ++k)
		for (int i = 1; i <= n; ++i) {
			parent[i][k + 1] = parent[parent[i][k]][k];
			cost[i][k + 1] = cost[parent[i][k]][k] + cost[i][k];
		}


	scanf("%d", &m);
	while (m--)
	{
		int sel;
		scanf("%d", &sel);

		if (sel == 1) {
			int from, to;
			scanf("%d %d", &from, &to);

			pair<int, long long> ret = LCA(from, to);

			printf("%lld\n", ret.second);


		}
		else
		{
			int from, to, k;
			scanf("%d %d %d", &from, &to, &k);

			pair<int, long long> ret = LCA(from, to);

			if (depth[from] - depth[ret.first] + 1 >= k) {

				int diff = k - 1;
				int cur = from;

				for (int i = 0; diff != 0; ++i) {
					if (diff % 2 == 1) {
						cur = parent[cur][i];
					}
					diff = diff >> 1;
				}

				printf("%d\n", cur);
			}
			else
			{
				k = k - (depth[from] - depth[ret.first]);

				int diff = depth[to] - depth[ret.first] + 1 - k;
				int cur = to;

				for (int i = 0; diff != 0; ++i) {
					if (diff % 2 == 1) {
						cur = parent[cur][i];
					}
					diff = diff >> 1;
				}

				printf("%d\n", cur);
			}

		}
	}
	return 0;
}

 

6. 결과