1761 정점들의 거리 - 백준

정점들의 거리

 

시간 제한메모리 제한제출정답맞힌 사람정답 비율

2 초

128 MB

11899

4761

3110

38.338%

문제

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

입력

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

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

출력

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

예제 입력 1 복사

7
1 6 13
6 3 9
3 5 7
4 1 3
2 4 20
4 7 2
3
1 6
1 4
2 6

예제 출력 1 복사

13
3
36

출처

• 문제를 번역한 사람: author6
• 문제의 오타를 찾은 사람: justin0907, wjsqjawns
• 빠진 조건을 찾은 사람: ntopia

풀이

N개의 정점으로 이루어진 트리가 주어지고 M개의 두 노드 쌍을 입력 받을 때 두 노드 사이 출력

 

입력

첫째 줄

N: 노드의 개수

N-1줄

두 점 및 거리를 받음

M 노드 쌍 갯수

M개의 줄

거리를 알고 싶은 노드 쌍이 한 쌍 씩 입력

 

정점은 1~N

 

M개의 줄에 차례로 입력받은 노드 거리 출력

 

DFS로 계산하면 정점*간선 길이 해서 40000*10000 = 4조 넘어가서 시간초과 예상

 

정점과 쿼리는 다 돌아아 함.

그러면 공통적으로 지나가는 경로들이 있는데 (트리 구조라서 조상 따라가는 형태) LCA 이용하면 될 것 같음

 

LCA는 첫주차였나에 O(MlogN)으로 풀수 있다고 배움.

 

 

위 예시를 트리로 표현

 

7
1 6 13
6 3 9
3 5 7
4 1 3
2 4 20
4 7 2
3
1 6
1 4
2 6

 

                                                    1

                                         /                                   \

                                6                                                       4

                         /                                                       /             \

              3                                                             2                      7

 

 

 

1-4의 사이 거리를 A라고 하면

dist[4] = A

4-7 사이의 거리를 B라고 하면

dist[7] = A+B

4-2 사이의 거리를 C라고 하면

dist[2] = A+ C

 

여기서 2-7 사이의 거리는 

B+C 이다. 즉, A를 2개 빼주면 된다.

 

1번 부터 n 번까지라고 했으니 루트를 1번으로 정한다.

트리의 루트에서 어떤 정점 까지의 거리를 dist[x]라고 한다.

 

그러면 정점 u부터 v까지의 길이는 위에서 계산한대로 dist[u] + dist[v] - 2*dist[lca(u,v)]를 풀어주면 된다.

 

코드

package test1761;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.*;
import java.util.StringTokenizer;

public class Main {

	static class Node {
		int end;
		int weight;

		public Node(int end, int weight) {
			this.end = end;
			this.weight = weight;
		}
	}

	static int N;
	static int M;
	static int treeHeight; // tree의 최대 높이
	static List<Node>[] list; // 입력받는 node list
	static int[] dist; // 루트 기준으로 노드 거리
	static int[] depth; // 해당 정점의 트리 높이
	static int[][] parent; // 해당 점점의 log(트리 높이)에 해당하는 부모 노드

	public static void main(String[] args) throws Exception {
		// TODO Auto-generated method stub
		MyScanner sc = new MyScanner(System.in);

		N = sc.nextInt();

		list = new ArrayList[N + 1];
		for (int i = 0; i < N + 1; i++) {
			list[i] = new ArrayList<>();
		}

		for (int i = 0; i < N - 1; i++) {
			int start = sc.nextInt();
			int end = sc.nextInt();
			int weight = sc.nextInt();

			list[start].add(new Node(end, weight));
			list[end].add(new Node(start, weight));
		}

		treeHeight = (int) Math.ceil(Math.log(N) / Math.log(2)) + 1;
		depth = new int[N + 1];
		dist = new int[N + 1];
		parent = new int[N + 1][treeHeight];

		init(1, 1, 0);

		// 나머지 2의 제곱승 번째의 부모 노드를 채워준다.
		for (int i = 1; i < treeHeight; i++) {
			for (int j = 1; j < N + 1; j++) {
				parent[j][i] = parent[parent[j][i - 1]][i - 1];
			}
		}

		M = sc.nextInt();
		
		for (int i = 0; i < M; i++) {
			int a = sc.nextInt();
			int b = sc.nextInt();

			// 최종 점화식
			sc.write(dist[a] + dist[b] - 2 * dist[algorithm(a, b)] + "\n");

		}
		sc.close();

	}

	// dfs 탐색을 통해 각 노드의 높이와 2^0번째 부모 노드 값으로 초기화
	static void init(int cur, int h, int pa) {
		depth[cur] = h;
		for (Node next : list[cur]) {
			if (next.end != pa) {
				dist[next.end] = dist[cur] + next.weight;
				init(next.end, h + 1, cur);
				parent[next.end][0] = cur;
			}
		}
	}

	// LCA 구하기
	static int algorithm(int a, int b) {
		int aHeight = depth[a];
		int bHeight = depth[b];

		// b의 depth가 a의 depth보다 높으면 swap
		if (aHeight < bHeight) {
			int temp = a;
			a = b;
			b = temp;
		}

		// 높이 맞추기
		for (int i = treeHeight - 1; i >= 0; i--) {
			if (Math.pow(2, i) <= depth[a] - depth[b]) {
				a = parent[a][i];
			}
		}

		if (a == b)
			return a;

		// LCA 찾기
		for (int i = treeHeight - 1; i >= 0; i--) {
			if (parent[a][i] != parent[b][i]) {
				a = parent[a][i];
				b = parent[b][i];
			}
		}
		return parent[a][0];
	}

	static class MyScanner {
		final BufferedReader reader;
		final BufferedWriter writer;
		static StringTokenizer tokenizer = null;

		MyScanner(InputStream in) {
			reader = new BufferedReader(new InputStreamReader(in));
			writer = new BufferedWriter(new OutputStreamWriter(System.out));
		}

		String nextToken() throws IOException {
			if (tokenizer == null || !tokenizer.hasMoreTokens()) {
				tokenizer = new StringTokenizer(reader.readLine());
			}

			return tokenizer.nextToken();
		}

		int nextInt() throws NumberFormatException, IOException {
			return Integer.parseInt(nextToken());
		}

		long nextLong() throws NumberFormatException, IOException {
			return Long.parseLong(nextToken());
		}

		void write(String sb) throws IOException {
			writer.write(sb);
		}

		void close() throws IOException {
			reader.close();
			writer.close();
		}
	}
}