[BOJ] 백준 10838번 트리 (Java)

#10838 트리

난이도 : 골드 1

유형 : 트리 / LCA

9466번: 텀 프로젝트

▸ 문제

N개의 노드로 구성된 루트가 있는 트리가 다음과 같이 주어진다. 각 노드는 0부터 N-1까지의 번호로 구별되고, 0번 노드는 루트 노드이고, 나머지 노드 각각은 0번 노드의 자식 노드이다.

트리에 적용할 수 있는 연산은 세 종류이며, 이를 통해 트리의 모양을 바꾸거나 트리 에지에 색칠을 할 수 있다. 각 연산과 그 의미는 다음과 같다. 

1. paint(a, b, c): a번 노드와 b번 노드를 잇는 최단 경로를 찾은 후, 경로 상에 있는 모든 에지를 색깔 c로 칠한다. 
2. move(a, b): a번 노드의 부모 노드를 b번 노드로 바꾼다. 단, b번 노드는 a번 노드를 루트로 하는 부트리(subtree)에 속하지 않는다. 부모 노드를 바꾸기 전 a번 노드의 부모 노드를 p라 할 때, 새로운 에지 (a,b)는 원래의 에지 (a,p)의 색깔을 갖는다. 
3. count(a, b): a번 노드와 b번 노드를 잇는 최단경로를 찾은 후, 그 경로 사이에 있는 에지에 칠해진 서로 다른 색깔의 개수를 출력한다. 

에지에 칠하는 색깔 c를 정수로 표시한다. 그리고 처음에는 모든 에지의 색깔이 0이라고 가정한다. 

예를 들어, 그림 1에서 보인 것처럼 6개의 노드로 구성된 초기 트리에 적용된 연산이 차례로

move(1,3); move(5,3); paint(5,4,8); move(3,4); paint(0,3,7); count(2,5);

일 때, 각 연산을 실행한 후 어떻게 트리의 모양과 에지 색깔이 바뀌는지를 아래 그림 2부터 그림 4에서 차례대로 보였다. 

그림 1. 초기 형태

 

그림 2. 좌측: move(1,3)을 실행한 후, 우측: move(5,3)을 실행 한 후

 

그림 3. paint(5,4,8)을 실행한 후

 

그림 4. 좌측: move(3,4)를 실행한 후, 우측: paint(0,3,7)을 실행한 후

 

그리고, 마지막 연산 count(2,5)에 대한 결과로는 3을 출력하게 된다. 왜냐하면, 그림 4의 우측 그림에서 보듯이 2번 노드와 5번 노드 사이의 최단 경로 상에 있는 에지들에 칠해진 색깔이 {0,7,8}로 3가지이기 때문이다. 

트리에 대한 정보와 일련의 연산이 주어질 때, 각 연산을 효과적으로 실행하는 프로그램을 작성하시오.

▸ 입력

첫째 줄에는 앞에서 설명한 트리의 노드 개수를 나타내는 정수 N(1 ≤ N ≤ 105)과 연산의 개수를 나타내는 정수 K(1 ≤ K ≤ 3×105)가 주어진다. 이어서 K 줄에 걸쳐 각 연산에 관한 정보가 한 줄에 하나씩 주어지는데, 각 줄에는 연산의 종류를 나타내는 정수 r(1 ≤ r ≤ 3)이 첫 번째로 주어진다.

• r = 1일 경우엔 연산이 paint 임을 의미하며, 세 정수 (a,b,c)가 추가로 같은 줄에 주어지는데, 여기서 a, b(0 ≤ a, b ≤ N-1)는 노드 번호를, c(0 ≤ c ≤109)는 색의 번호를 나타낸다.
• r = 2일 경우엔 연산이 move임을 의미하며, 두 정수 a, b(1 ≤ a ≤ N-1, 0 ≤ b ≤ N-1)가 추가로 같은 줄에 주어지는데, 이는 노드 번호를 나타낸다. 
• r = 3일 경우엔 연산이 count임을 의미하며, 두 정수 a, b(0 ≤ a, b ≤ N-1)가 추가로 같은 줄에 주어지는데, 이는 노드 번호를 나타낸다. 

노드의 개수가 N인 트리의 초기 모양은 그림 1에서 보인 것처럼 0번 노드가 루트이고, 나머지 노드들의 부모 노드는 0번 노드이며, 초기 트리의 모든 에지 색깔은 0이라고 가정한 사실을 기억하기 바란다. 

또한, paint와 count 연산 시 a번 노드와 b번 노드 사이의 최단경로의 길이는 항상 1,000 이하이다.

▸ 출력

입력에서 주어진 count 연산 각각에 대해, 그 순서대로 그 때의 결과 값을 한 줄에 출력한다. 

 

문제 풀이  

LCA를 활용한 트리 문제이다. 원래 DP를 사용해서 부모노드를 2^h만큼 점프 탐색을 하면서 O(N*M) 시간복잡도를 O(MlogN)으로 축소시켜줘야 하지만 count 함수에서 부모노드를 전부 탐색해줘야 하기 떄문에 일일히 모든 부모노드를 탐색해주는 브루트포스한 방법을 사용해야 한다. 그래서 문제에서도 시간 초과가 발생하지 않기 위해 '또한, paint와 count 연산 시 a번 노드와 b번 노드 사이의 최단경로의 길이는 항상 1,000 이하이다.' 라는 조건을 추가해줬다.

 

LCA 구하기

1. a와 b 노드가 같다면 LCA는 a == b이므로 해당 노드를 바로 출력한다.
2. a의 부모노드를 먼저 탐색하고 해당 경로를 체크해준다. 최대 길이는 1000이기 떄문에 반복문을 1000번으로 제한해둔다.
3. b의 부모노드를 탐색하면서 만약 탐색 여부가 체크(a의 부모노드)가 되어있다면 해당 노드가 LCA이므로 바로 출력해준다.
4. 모든 탐색 후에도 공통된 부모가 없다면 루트노드인 0을 출력한다.

static int findLCA(int a, int b) {
	if(a==b) return a;
	
	boolean[] check = new boolean[n];
	for(int i=0; i<MAX_LEN; i++) { // a 노드 부모 탐색 
		check[a] = true;
		a = nodes[a].pa;
	}
	
	for(int i=0; i<MAX_LEN; i++) { // b 노드 부모 탐색
		if(check[b]) return b;
		b = nodes[b].pa;
	}
	return 0;
}

 

풀이 코드 

import java.io.*;
import java.util.*;

public class Main {

	static class Node{
		int pa;
		int color;
		public Node(int pa, int color) {
			this.pa = pa;
			this.color = color;
		}
	}
	static final int MAX_LEN = 1000;
	static int n;
	static Node[] nodes;
	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 k = Integer.parseInt(st.nextToken());
		nodes = new Node[n];
		for(int i=0; i<n; i++) {
			nodes[i] = new Node(0,0);
		}
		
		StringBuilder sb = new StringBuilder();
		for(int i=0; i<k; i++) {
			st = new StringTokenizer(br.readLine());
			int r = Integer.parseInt(st.nextToken());
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			
			if(r==1) {
				int c = Integer.parseInt(st.nextToken());
				paint(a,b,c);
			}else if(r==2) {
				move(a,b);
			}else {
				sb.append(count(a,b)+"\n");
			}
		}
		System.out.println(sb.toString());
	}
	
	static void move(int a, int b) {
		nodes[a].pa = b;
	}
	
	static int findLCA(int a, int b) {
		if(a==b) return a;
		
		boolean[] check = new boolean[n];
		for(int i=0; i<MAX_LEN; i++) { // a 노드 부모 탐색 
			check[a] = true;
			a = nodes[a].pa;
		}
		
		for(int i=0; i<MAX_LEN; i++) { // b 노드 부모 탐색
			if(check[b]) return b;
			b = nodes[b].pa;
		}
		return 0;
	}
	
	static void paint(int a, int b, int c) {
		int lca = findLCA(a,b);
		while(a != lca) {
			nodes[a].color = c;
			a = nodes[a].pa;
		}
		
		while(b != lca) {
			nodes[b].color = c;
			b = nodes[b].pa;
		}
	}
	
	static int count(int a, int b) {
		int lca = findLCA(a,b);
		Set<Integer> set = new HashSet<>();
		while(a != lca) {
			set.add(nodes[a].color);
			a = nodes[a].pa;
		}
		
		while(b != lca) {
			set.add(nodes[b].color);
			b = nodes[b].pa;
		}
		return set.size();
	}
}