https://www.acmicpc.net/problem/15681

 

15681번: 트리와 쿼리

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105) 이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

www.acmicpc.net

코드설명

트리 DFS 트리의 DP를 활용하여 해결하는 문제입니다.

 

처음에는 R을 루트로 하는 트리를 아래와 같이 새로 만들어서 DFS를 통해 진행하려고했습니다만,

//5를 root로 하는 새로운 트리 makeTree 생성됨.
public static void makeTree(int currentNode, int idx) {
	
	for(int i=0;i<graph.get(currentNode).size();i++) {
		int temp = graph.get(currentNode).get(i);
		if(temp != idx) {
			newTree.get(currentNode).add(temp);
			makeTree(temp, currentNode);
		}
	}
}

입력조건에

- 입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

- 임의의 두 정점 U와 V에 대해, U에서 V로 가는 최단경로는 유일하다.

- 아무 정점이나 잡고 부모와의 연결을 끊었을 때, 해당 정점과 그 자식들, 그 자식들의 자식들… 로 이루어진 부분그래프는 트리가 된다.

- 트리에서는 (눈치챘을 수도 있지만) 어떤 정점의 부모는 하나이거나 없다

즉, 트리에서 아무 노드를 루트노드로 잡고 트리라고 판단해도 문제가 없습니다. 

public static void treeDFSAndDP(int currentNode, int beforeNode) {
    treeDp[currentNode] = 1; // 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.

    for(int i=0;i<graph.get(currentNode).size();i++) {
        int temp = graph.get(currentNode).get(i);
        if( temp != beforeNode) {
            treeDFSAndDP(temp, currentNode);
            treeDp[currentNode] += treeDp[temp];
        }

    }
}

- currentNode는 현재 탐색중인 노드입니다.

- temp는 새로 탐색할 노드입니다.

 

특히 이문제에서 가장 유의해야할점은 treeDp[currentNode] += treeDp[temp] 를 통해 탐색이 끝난 DFS가 저장한 트리의 자식들의 개수를 더해준다는 점입니다. 

해당 논리가 이 문제에서 가장 중요한 점이라고 생각합니다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {
	public static int N, R, Q;
	public static int[] arr;
	public static ArrayList<ArrayList<Integer>> graph = new ArrayList<ArrayList<Integer>>();
	public static int[] treeDp;
	public static int answer = 0;
    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());
    	R = Integer.parseInt(st.nextToken());
    	Q = Integer.parseInt(st.nextToken());
    	
    	for(int i=0;i<=N;i++) {
    		graph.add(new ArrayList<>());
    	}
    	treeDp = new int[N+1];
    	
    	for(int i=0;i<N-1;i++) {
    		st = new StringTokenizer(br.readLine());
    		int a = Integer.parseInt(st.nextToken());
    		int b = Integer.parseInt(st.nextToken());
    		graph.get(a).add(b);
    		graph.get(b).add(a);
    	}
    	
    	treeDFSAndDP(R, -1);
    	
    	for(int i=0;i<Q;i++) {
    		st = new StringTokenizer(br.readLine());
    		int U = Integer.parseInt(st.nextToken());
    		System.out.println(treeDp[U]);
    	}
    	
    	
    	
    }
    
    public static void treeDFSAndDP(int currentNode, int beforeNode) {
    	treeDp[currentNode] = 1; // 자신도 자신을 루트로 하는 서브트리에 포함되므로 0이 아닌 1에서 시작한다.
    	
    	for(int i=0;i<graph.get(currentNode).size();i++) {
    		int temp = graph.get(currentNode).get(i);
    		if( temp != beforeNode) {
    			treeDFSAndDP(temp, currentNode);
    			treeDp[currentNode] += treeDp[temp];
    		}
    		
    	}
    }
        
}

+ Recent posts