[C++]백준 - 13511번 문제

13511번: 트리와 쿼리 2 (acmicpc.net)

13511번: 트리와 쿼리 2

 

13511번 : 트리와 쿼리 2

---

N개의 정점으로 이루어진 트리(무방향 사이클이 없는 연결 그래프)가 있다. 정점은 1번부터 N번까지 번호가 매겨져 있고, 간선은 1번부터 N-1번까지 번호가 매겨져 있다.

아래의 두 쿼리를 수행하는 프로그램을 작성하시오.

• 1 u v: u에서 v로 가는 경로의 비용을 출력한다.
• 2 u v k: u에서 v로 가는 경로에 존재하는 정점 중에서 k번째 정점을 출력한다. k는 u에서 v로 가는 경로에 포함된 정점의 수보다 작거나 같다.

 

입력

---

첫째 줄에 N (2 ≤ N ≤ 100,000)이 주어진다.

둘째 줄부터 N-1개의 줄에는 i번 간선이 연결하는 두 정점 번호 u와 v와 간선의 비용 w가 주어진다.

다음 줄에는 쿼리의 개수 M (1 ≤ M ≤ 100,000)이 주어진다.

다음 M개의 줄에는 쿼리가 한 줄에 하나씩 주어진다.

간선의 비용은 항상 1,000,000보다 작거나 같은 자연수이다.

 

 

출력

---

각각의 쿼리의 결과를 순서대로 한 줄에 하나씩 출력한다.

 

---

 

생각해 볼 점

---

11438번 문제와 거의 동일합니다.

 

11438 풀이에서 간선의 W(가중치)만 추가로 더해주면 금방 풀립니다.

 

코드

---

#include <iostream>
#include <cmath>
#include <vector>
using namespace std;

int N, logN;
int* level;
pair<int, long long>** sparseTable;
vector<pair<int, int>>* edges;	//<node, cost>

//lv만큼 위에 있는 parent 노드를 구함
pair<int, long long> getParent(int node, int lv)
{
	long long cost = 0;
	while (lv)
	{
		int log_diff = log2(lv);
		lv -= 1 << log_diff;

		cost += sparseTable[node][log_diff].second;
		node = sparseTable[node][log_diff].first;
	}

	return { node, cost };
}

//레벨(깊이)를 세팅하는 함수
void setLevels(int const& node)
{
	for (pair<int, int> nextNode : edges[node])
	{
		if (level[nextNode.first] != -1) continue;
		level[nextNode.first] = level[node] + 1;

		sparseTable[nextNode.first][0] = { node, nextNode.second };
		setLevels(nextNode.first);
	}
}

//희소 행렬 세팅하는 함수
void setSparseTable()
{
	for (int i = 1; i < logN; i++)
	{
		for (int node = 0; node < N; node++)
		{
			int newNode = sparseTable[node][i - 1].first;
			
			if (newNode == -1) continue;
			long long newCost = sparseTable[node][i - 1].second + sparseTable[newNode][i-1].second;
			newNode = sparseTable[newNode][i - 1].first;
			sparseTable[node][i] = { newNode, newCost };
		}
	}
}

int main()
{
	const pair<int, long long> defPair = { -1, 0 };

	scanf("%d", &N);

	logN = log2(N) + 1;
	level = new int[N];
	sparseTable = new pair<int, long long>*[N];
	edges = new vector<pair<int, int>>[N];

	//간선 입력 및 초기화
	for (int i = 0; i < N - 1; i++)
	{
		int a, b, w;
		scanf("%d %d %d", &a, &b, &w);

		a--; b--;
		edges[a].push_back({ b, w });
		edges[b].push_back({ a, w });

		level[i + 1] = -1;
		sparseTable[i] = new pair<int, long long>[logN];
		std::fill_n(sparseTable[i], logN, defPair);
	}
	sparseTable[N - 1] = new pair<int, long long>[logN];
	std::fill_n(sparseTable[N - 1], logN, defPair);

	level[0] = 0;
	setLevels(0);
	setSparseTable();

	int M;
	scanf("%d", &M);
	

	//쿼리
	for (int i = 0; i < M; i++)
	{
		int mode;
		scanf("%d", &mode);
		
		int a, b;
		long long ans = 0;

		scanf("%d %d", &a, &b);
		a--; b--;

		int originalA = a, originalB = b;

		//swap
		if (level[b] < level[a])
		{
			int temp = a;
			a = b;
			b = temp;
		}

		int lv_diff = level[b] - level[a];
		
		//레벨 맞춰주기
		pair<int, long long> ret = getParent(b, lv_diff);
		b = ret.first;
		ans += ret.second;

		//공통 부모 구하기 + 가중치 덧셈
		while (a != b)
		{
			int log_lv = log2(level[a]);
			for (; 0 < log_lv; log_lv--)
			{
				if (sparseTable[a][log_lv].first != sparseTable[b][log_lv].first) break;
			}

			ans += sparseTable[a][log_lv].second;
			ans += sparseTable[b][log_lv].second;
			a = sparseTable[a][log_lv].first;
			b = sparseTable[b][log_lv].first;
		}
		int commonParent = a;

		if (mode == 1) printf("%lld\n", ans);
		//K번째 노드 구하기
		else if(mode == 2)
		{
			int k;
			scanf("%d", &k);
			k--;
			//k가 A -> 최대 공통 부모까지의 거리보다 클 경우 B에서 구함
			int from = originalA;
			if (k > level[originalA] - level[commonParent])
			{
				k = level[originalA] + level[originalB] - (level[commonParent] << 1) - k;
				from = originalB;
			}
			printf("%d\n", getParent(from, k).first + 1);
		}
	}


	for (int i = 0; i < N; i++) delete[] sparseTable[i];
	delete[] level;
	delete[] sparseTable;
	delete[] edges;
	return 0;
}

 

그 외

---