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