[백준] 15681번 - "트리와 쿼리" (Java)

www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

 

---

 

난이도 : 골5

 

트리 연습하기 좋은 문제라고 생각합니다. 트리 만드는 법을 배울 수 있습니다. 그리고 

이 문제도 테스트 케이스 문제라서 각 테스트 케이스마다 문제를 해결하고자 하면 시간 초과가 날 수 있습니다. 미리 문제를 다 풀고 테스트 케이스 별로 결과만 출력하는 것을 추천합니다.

 

이 문제는 힌트가 아주 잘 나와있어 힌트에 나와있는 의사코드(pseudocode)를 따라서 작성하면 쉽게 풀 수 있습니다.

 

- void main()  

모든 입력과 전역 변수들을 선언 및 초기화 합니다. makeTree()함수를 호출하여 트리를 만들고

countSubtreeNodes()를 호출하여 재귀 및 dp를 활용하여 subTree의 갯수를 구합니다.

 

- void makeTree(int curNode, int p)

list 변수를 바탕으로 Tree(subTree)를 구성합니다. 매개 변수의 curNode는 현재 노드를 나타냅니다.

p는 부모 노드를 의미합니다. 즉 main 문에서 makeTree(r, -1)을 호출하면 루트 노드를 기준으로 트리를 만듭니다. 루트 노드는 부모가 없기 때문에 매개 변수 p에 -1 값을 넣어줍니다. 

 

부모 노드의 정보를 가져가면서, 부모 노드가 아니면서 자신과 연결되어 있는 모든 노드를 자신의 자식으로, 자신의 자식이 될 노드들의 부모 노드를 자신으로 연결한 뒤 재귀적으로 자식 노드들에게 트리 구성을 요청하는 형태의 함수입니다.

 

- void countSubtreeNodes(int curNode, int p)

재귀 호출과 dp를 활용하여 서브 트리에 포함되어 있는 모든 노드들의 갯수를 구합니다. 

자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작합니다.

 

자식 노드들에 대한 makeTree()가 호출된 뒤엔, 해당 자식 노드를 서브트리로 하는 트리가 구성이 완료됩니다. 이와 같은 원리로 모든 노드에 대해 해당 노드를 루트로 하는 서브트리에 속한 노드의 수를 계산하는 함수입니다.

 

메인 함수에서 countSubtreeNodes(r)를 호출하게 되면, r번을 루트로 하는 트리에서 모든 노드에 대해 각 노드를 루트로 하는 서브트리에 속한 정점의 수를 계산해둘 수가 있습니다. 즉, 쿼리를 호출하기 전에 각 노드의  서브 트리에 속한 모든 노드들의 갯수를 미리 구해 놓을 수 있습니다. 이를 활용하여, 주어진 시간 제한 내에서 문제를 충분히 풀 수 있게됩니다.

 

 

모든 함수들이 다 하나의 for문 안에서 동작하므로 (재귀 호출도 N만큼 호출됩니다.)

최악의 시간 복잡도는 앞의 상수는 생략하므로 O(N)입니다.

 

  코드

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

public class BOJ_15681 {
	// BOJ / 15681번 / 트리와 쿼리 / 트리 / 골드5
	static ArrayList<Integer>[] list, tree;
	static int parent[], size[];
	public static void main(String[] args) throws IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer stk = new StringTokenizer(br.readLine(), " ");
		
		int n = Integer.parseInt(stk.nextToken());
		int r = Integer.parseInt(stk.nextToken());
		int q = Integer.parseInt(stk.nextToken());
		
		parent = new int [n+1];
		size = new int[n+1];
		list = new ArrayList[n+1];
		tree = new ArrayList[n+1];
		
		for(int i=0; i<list.length; i++) {
			list[i] = new ArrayList<>();
			tree[i] = new ArrayList<>();
		}
		
		for(int i=1; i<n; i++) {
			stk = new StringTokenizer(br.readLine(), " ");
			int u = Integer.parseInt(stk.nextToken());
			int v = Integer.parseInt(stk.nextToken());
			
			list[u].add(v);
			list[v].add(u);
		}
		
		makeTree(r, -1);
		countSubtreeNodes(r);
		StringBuffer sb = new StringBuffer();
		while(q-- > 0) {
			int query = Integer.parseInt(br.readLine());
			
			sb.append(size[query]).append("\n");
		}
		System.out.println(sb);
	}
	
	// list를 바탕으로 tree(Subtree) 구성한다. 
	// p가 -1은 부모가 없음을 의미한다. 즉 루트 노드이다.
	public static void makeTree(int curNode, int p) {
		
		for(int node : list[curNode]) {
			if(node != p) {
				tree[curNode].add(node);
				parent[node] = curNode;
				makeTree(node, curNode);
			}
		}
	}
	
	// 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.
	public static void countSubtreeNodes(int curNode) {
		size[curNode] = 1;
		
		for(int node : tree[curNode]) {
			countSubtreeNodes(node);
			size[curNode] += size[node];
		}
	}
}

 

추가적으로 문제의 힌트에 나온 방식이 아닌 코드도 첨부하겠습니다. 

위의 코드보다는 조금(?) 더 빠릅니다. tree를 만들지 않고 기본 list 변수만 생성하고 

그 list를 활용하여 dfs 함수를 호출함과 동시에 dp에 값을 저장하는 코드입니다. 

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

public class BOJ_15681 {
	// BOJ / 15681번 / 트리와 쿼리 / 트리 / 골드5
	static ArrayList<Integer>[] list;
	static boolean visit[];
	static int dp[];
	public static void main(String[] args) throws NumberFormatException, IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer stk = new StringTokenizer(br.readLine(), " ");
		
		int n = Integer.parseInt(stk.nextToken());
		int r = Integer.parseInt(stk.nextToken());
		int q = Integer.parseInt(stk.nextToken());
		
		list = new ArrayList[n+1];
		dp = new int [n+1];
		visit = new boolean[n+1];
		
		for(int i=0; i<list.length; i++) {
			list[i] = new ArrayList<>();
		}
		
		for(int i=1; i<n; i++) {
			stk = new StringTokenizer(br.readLine(), " ");
			int u = Integer.parseInt(stk.nextToken());
			int v = Integer.parseInt(stk.nextToken());
			
			list[u].add(v);
			list[v].add(u);
		}
		
		dfs(r);
		StringBuffer sb = new StringBuffer();
		while(q-- > 0) {
			int query = Integer.parseInt(br.readLine());
			
			sb.append(dp[query]).append("\n");
		}
		System.out.println(sb);
	}
	
	public static int dfs(int now) {
		
		if(dp[now] != 0) return dp[now];
		visit[now] = true;
		int count = 1;
		
		for(int node : list[now]) {
			if(visit[node]) continue;
			count += dfs(node);
		}
		dp[now] = count;
		
		return dp[now];
	}
}

 이건 여담이지만, 테스트 케이스 문제를 풀 때는 무조건 StringBuilder() 혹은 StringBuffer()를 사용하여 출력하는 것을 

추천드립니다. 시간이 차이 나는 것은 알았지만 이렇게 차이가 많이 날 줄은 몰랐습니다. 

 

1000ms 이하로 나온 코드가 StringBuffer()를 이용하여 출력한 코드입니다. 

 

 

 

 

피드백은 얼마든지 환영합니다.