15681 트리와 쿼리

DONGJIN IM·2022년 3월 8일
0

코딩 테스트

목록 보기
66/137

문제 이해

트리가 주어진다. 또한 Root로 지정할 Node도 주어질 것이다.
트리 구성 요소가 주어진 이후, N개의 번호가 주어질 것이다. 이 번호를 U라고 하자.

이 때, U를 루트로 하는 서브트리가 가진 Node 개수를 출력하는 문제이다.
(서브트리이므로, U도 Node 개수에 포함되어 있어야 한다)


문제 풀이

A의 왼쪽 자식을 B, 오른쪽 자식을 C라고 생각하자.
이 때 A를 Root로 가지는 Subtree의 node 개수는 몇 개일까?
이는 (B를 Root로 가지는 SubTree의 node 개수) + (C를 Root로 가지는 SubTree의 node 개수) + 1(A 자기 자신)이라고 할 수 있다.
이는 자식이 2개가 아니라 매우 많아져도 똑같을 것이다.

이 방법을 활용하여 문제를 풀자.

먼저 Leaf Node의 SubTree의 Node 개수는 1이 될 것이다.
그리고 Leaf Node의 부모 Node는 (다른 자식 Node의 SubTree node 개수) + 1이 될 것이다. 즉, Leaf Node의 계산 값이 부모 Node에 영향을 주는 방식이다.

따라서 Leaf Node부터 검사하여 값을 변경하고, 이를 부모 Node에 전파해주는 알고리즘을 짜야하며, 이에 가장 적절한 DFS를 활용했다.(DFS는 Leaf Node에 도달해야지만 재귀함수가 종료되기 때문)

  1. Leaf Node임 : 해당 SubTree 개수를 1로 지정 이후 1을 반환

  2. Leaf Node가 아님 : 모든 자식 Node에 대하여 DFS 수행. 이후 DFS를 통해 반환된 값들을 모두 합치고, 자기 자신을 합쳐(즉, 1을 더해) SubTree 개수를 지정하고 이렇게 구한 SubTree 개수를 반환함

이후, 배열에 저장한 SubTree 개수를 arr[index]를 통해 바로 접근하여 답을 구할 수 있을 것이다.


코드

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

public class Main {
	static StringBuilder sb = new StringBuilder();
	static int N, R, Q;
	static ArrayList<Integer>[] nodes;
	static boolean[] visit;
	static int[] answer;
	
	static Integer dfs(int start) {
		visit[start] = true;
		
		int ans = 1; 
        // SubTree Node 개수에는 Node start도 포함되어 있어야 하므로 
        // 1로 초기화 시켜준다.
		for(int s:nodes[start]) {
			if(!visit[s]) {
                // visit[s] = true라면 이미 방문했다는 의미
                // 이미 방문했다는 것은 부모 Node라는 의미이므로 무시한다.
				ans+=dfs(s);
			}
		}
		
		answer[start] = ans; 
        // start를 Root로 하는 SubTree의 Node개수를 배열에 저장 
		
		return ans; 
        // start를 Root로 하는 SubTree의 Node 개수 반환.
        // 이 과정이 이뤄져야 ans+=dfs(s)를 통해 
        // SubTree의 전체 Node개수를 구할 수 있음
	}
	
	public static void main(String[] args) throws IOException {
		FastReader sc = new FastReader();
		
		N = sc.nextInt();
		R = sc.nextInt(); 
		Q = sc.nextInt(); 
		
		nodes = new ArrayList[N+1];
		visit = new boolean[N+1];
		answer = new int[N+1];
		for(int i =1;i<N+1;i++) nodes[i] = new ArrayList<>();
		
		for(int i = 0;i<N-1;i++) {
			int tmp1 = sc.nextInt();
			int tmp2 = sc.nextInt();
			
			nodes[tmp1].add(tmp2);
			nodes[tmp2].add(tmp1);
		}
		
		dfs(R); // R이 Root Node이므로 R을 시작으로 DFS를 수행해야 한다.
		
		for(int i =0;i<Q;i++) {
			int tmp1 = sc.nextInt();
			sb.append(answer[tmp1]).append("\n");
		}
		System.out.println(sb);
	}
	
	static class FastReader // 빠른 입력을 위한 클래스
}

결과

profile
개념부터 확실히!

0개의 댓글

관련 채용 정보