Boj 13511) 트리와 쿼리 2

문제

백준 13511

설명

희소배열을 사용한 LCA 로 로그시간에 공통조상을 찾아서 u, v 와의 관계로 쿼리를 풀어내면 되는 문제이다.

Heavy Light Decomposition 으로도 풀리지만 정해는 아니라는듯.

시간 복잡도

O(\((\mathrm{N} + \mathrm{M})\log{\mathrm{N}}\))

코드

using ll = long long;
struct E { int d, w; };
struct Node {
	vector<E> childs;
	int depth = 0;
};

#define MAX_SIZE 100001
Node nodes[MAX_SIZE];
int  parentTB[20][MAX_SIZE];
ll   weights[MAX_SIZE];

void DFS(int cur, int parent = 0, ll weight = 0)
{
	nodes[cur].depth = nodes[parent].depth + 1;
	parentTB[0][cur] = parent;
	weights[cur] = weight;
	for (auto child : nodes[cur].childs)
		if (parent != child.d)
			DFS(child.d, cur, weight + child.w);
}

int LCA(int a, int b)
{
	int i;
	if (nodes[a].depth < nodes[b].depth) swap(a, b);
	for (int dff = nodes[a].depth - nodes[b].depth; dff > 0; )
	{
		for (i = 0; 2 << i <= dff; i++);
		a = parentTB[i][a];
		dff = nodes[a].depth - nodes[b].depth;
	}

	while (a != b)
	{
		for (i = 1; parentTB[i][a] != parentTB[i][b]; i++);
		a = parentTB[i - 1][a];
		b = parentTB[i - 1][b];
	}
	
	return a;
}

int LCA(int a, int b, int k)
{
	k--; // 첫번째가 한칸 뒤가 아니라 시작노드라 단위를 맞춤
	
	int root = LCA(a, b);
	int dist = nodes[a].depth + nodes[b].depth - (nodes[root].depth << 1);
	
	// a 에서 b 로의 경로의 k 번째 노드가 LCA 를 거쳐간다면 b 에서부터 부모를 찾아야함
	if (nodes[a].depth - nodes[root].depth < k)
	{
		k = (nodes[b].depth - nodes[root].depth) - (k - (nodes[a].depth - nodes[root].depth));
		swap(a, b);  // 편의상 b 를 a 로 바꿈
	}
	
	// a 로부터 k 번째 부모를 찾음
	while (k) {
		int i;
		for (i = 0; (2 << i) <= k; i++);
		a = parentTB[i][a];
		k -= 1 << i;
	}
	return a;
}

int main()
{
	int n, q;
	cin >> n;
	for (int i = 0; i < n - 1; i++)
	{
		int a, b, c;
		cin >> a >> b >> c;
		nodes[a].childs.push_back({ b, c });
		nodes[b].childs.push_back({ a, c });
	}
	
	DFS(1, 0);
	for (int i = 1; i < 20; i++)
		for (int j = 1; j <= n; j++)
			parentTB[i][j] = parentTB[i - 1][parentTB[i - 1][j]];
	
	cin >> q;
	while (q--)
	{
		int type;
		cin >> type;
		if (type == 1)
		{
			// from u -> v weight
			int u, v;
			cin >> u >> v;
			cout << weights[u] + weights[v] - (weights[LCA(u, v)] << 1) << '\n';
		}
		else {
			// from u -> v kth vertex, kth exist
			int u, v, w;
			cin >> u >> v >> w;
			cout << LCA(u, v, w) << '\n';
		}
	}
}