[BOJ] 백준 1240번 노드사이의 거리 (Java)

#1240 노드사이의 거리

난이도 : 골드 5

유형 : 트리 / DFS

1240번: 노드사이의 거리

▸ 문제

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

▸ 입력

첫째 줄에 노드의 개수 N이 입력되고 다음 N-1개의 줄에 트리 상에 연결된 두 점과 거리(10,000 이하의 정수)를 입력받는다. 그 다음 줄에는 거리를 알고 싶은 M개의 노드 쌍이 한 줄에 한 쌍씩 입력된다.

▸ 출력

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

 

문제 풀이  

트리의 크기가 크게 주어졌다면 LCA 알고리즘을 사용하여 쿼리 최적화를 해주면되지만 해당 문제는 쿼리의 수랑 트리 크기가 작으므로 그냥 DFS탐색으로 두 노드 사이의 거리를 구해주면 된다.

 

설계

1. 노드 데이터와 거리 비용을 인접리스트 배열에 저장한다. list[u].add(new Node(v,w));
2. dfs탐색을 통해 to노드부터 from노드까지의 이동 거리를 구한다.
   1. if(to == from) answer = cost;

 

풀이 코드 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

public class Main {

	static class Node{
		int to;
		int cost;
		
		public Node(int to, int cost) {
			this.to = to;
			this.cost = cost;
		}
	}
	static int answer;
	static List<Node>[] list;
	static int[] cost;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		
		int n = Integer.parseInt(st.nextToken());
		int m = Integer.parseInt(st.nextToken());
		
		list = new ArrayList[n+1];
		cost = new int[n+1];
		for(int i=1; i<n+1; i++) {
			list[i] = new ArrayList<>();
		}
		
		for(int i=0; i<n-1; i++) {
			st = new StringTokenizer(br.readLine());
			int u = Integer.parseInt(st.nextToken());
			int v = Integer.parseInt(st.nextToken());
			int w = Integer.parseInt(st.nextToken());
			list[u].add(new Node(v,w));
			list[v].add(new Node(u,w));
		}
        
		for(int i=0; i<m; i++) {
			st = new StringTokenizer(br.readLine());
			int to = Integer.parseInt(st.nextToken());
			int from = Integer.parseInt(st.nextToken());
			dfs(to, -1, from, 0);
			System.out.println(answer);
		}
	}
	
	static void dfs(int to, int pa, int from, int cost) {
		if(to == from) {
			answer = cost;
		}
		for(Node nxt : list[to]) {
			if(nxt.to != pa) {
				dfs(nxt.to, to, from, cost+nxt.cost);
			}
		}
	}
}

 

 

LCA 풀이

백준 1761번 정점들의 거리는 이와 같은 문제인데, 노드 최대 크기는 4만개, 쿼리는 최대 1만개로 위의 알고리즘으로 해결할 수 없다. 그래서  LCA알고리즘을 통해 풀면 다음과 같다.

import java.io.*;
import java.util.ArrayList;
import java.util.List;
import java.util.StringTokenizer;

public class Main {

	static class Node{
		int to;
		int w;
		
		public Node(int to, int w) {
			this.to = to;
			this.w = w;
		}
	}
	
	static int n,h;
	static List<Node>[] list;
	static int[][] dp;
	static int[] dis;
	static int[] depth;
	static StringBuilder sb = new StringBuilder();
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		n = Integer.parseInt(st.nextToken());
		int m = Integer.parseInt(st.nextToken());
		
		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++) {
			st = new StringTokenizer(br.readLine());
			int from = Integer.parseInt(st.nextToken());
			int to = Integer.parseInt(st.nextToken());
			int w= Integer.parseInt(st.nextToken());
			
			list[from].add(new Node(to,w));
			list[to].add(new Node(from,w));
		}
		
		h = getTreeH();
		depth = new int[n+1];
		dis = new int[n+1];
		dp = new int[n+1][h];
		
		init(1,1,0);
		fillParents();
		for(int i=0; i<m; i++) {
			st = new StringTokenizer(br.readLine());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			int res = LCA(a,b);
			sb.append(dis[a] + dis[b] -2*dis[res]).append("\n");
			
		}
		System.out.println(sb.toString());
	}
	
	static int getTreeH() {
		return (int)Math.ceil(Math.log(n)/Math.log(2))+1;
	}
	
	static void init(int cur, int h, int pa) {
		depth[cur] = h;
		for(Node nxt : list[cur]) {
			if(nxt.to!=pa) {
				dis[nxt.to] = dis[cur] + nxt.w;
				init(nxt.to, h+1, cur);
				dp[nxt.to][0] = cur;
			}
		}
	}
	
	static void fillParents() {
		for(int i=1; i<h; i++) {
			for(int j=1; j<n+1; j++) {
				dp[j][i] = dp[dp[j][i-1]][i-1];
			}
		}
	}
	
	static int LCA(int a, int b) {
		int ah = depth[a];
		int bh = depth[b];
		
		if(ah < bh) {
			int tmp =a;
			a= b;
			b= tmp;
		}
		
		for(int i=h-1; i>=0; i--) {
			if(Math.pow(2, i) <= depth[a] - depth[b]) {
				a = dp[a][i];
			}
		}
		
		if(a==b) return a;
		
		for(int i=h-1; i>=0; i--) {
			if(dp[a][i] != dp[b][i]) {
				a = dp[a][i];
				b = dp[b][i];
			}
		}
		return dp[a][0];
	}
}

 

 

 LCA알고리즘을 사용한 코드(위), DFS을 사용한 코드(아래)의 채점 결과는 다음과 같다.

 

채점 결과