[DFS+DP] 백준 15681번 트리와 쿼리 - JAVA

 

 

 

오늘의 문제는 트리문제이다.

사실 간단한 트리문제인 줄 알고 접근했다가 "메모리 초과"를 당해버렸다 ㅎㅎ...

 

---

1. 출처

---

 

15681번: 트리와 쿼리

 

 

 

---

2. 설계

---

 

처음에는 쿼리수가 하나씩 주어지면 그 때마다 서브트리를 탐색하면서 count를 세줬다. 이 때 BFS를 사용했다. Queue를 사용하다보니 메모리 초과가 난 것이다. 큐를 안쓰고 어떻게 풀까? 고민했다.

 

그래서 모범 답안을 찾아보니 DFS + DP로 푼 개발자 분들이 많았다.

루트 노트부터 시작해 DFS를 이용해 모든 자식들을 탐색하고, 더 이상 탐색할 자식이 없을 때마다 DP를 이용해 자식 수를 저장하면 된다.

 

 

1. 무향 그래프 입력 받기

for(int i=1;i<=N-1;i++){
	String UV = br.readLine();
	st = new StringTokenizer(UV, " ");
	int U = Integer.parseInt(st.nextToken());
	int V = Integer.parseInt(st.nextToken());

	adj[U].add(V);
	adj[V].add(U);
}

 

 

 

2. DFS + DP 자식 수 세기

public static void SubTree(int child, int parent){

	for(int i=0;i<adj[child].size();i++){
		int next = adj[child].get(i);
		if(next == parent) continue;
		SubTree(next, child);
	}

	if(parent!=-1){//부모가 있다면
		childCount[parent] += childCount[child];
	}
}

 

 

 

예를 들어 이렇게 for문을 다 돌 때마다 자식 수를 부모에 dp로 넣어주는 것이다.

 

 

 

 

대략 이런 식으로 전체 트리 노드의 자식 수를 세준다. 중간 과정일 뿐이지만 노드 4의 자식이 자신을 포함해 4개라는 답을 얻을 수 있다.

 

 

 

3. 답 출력

sb = new StringBuilder();
for(int T=1;T<=Q;T++){
	int U = Integer.parseInt(br.readLine());
	sb.append(childCount[U]).append("\n");
}

System.out.println(sb);

 

 

 

---

3. 전체 코드

---

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

class Main {
    
    static List<Integer>[] adj;
    static int[] childCount;
    static StringBuilder sb;
    
    public static void main(String[] args) throws IOException {

        //1. 입력
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        String NRQ = br.readLine();
        StringTokenizer st = new StringTokenizer(NRQ, " ");
        int N = Integer.parseInt(st.nextToken());
        int R = Integer.parseInt(st.nextToken());
        int Q = Integer.parseInt(st.nextToken());

        adj = new ArrayList[N+1]; //정점 1부터 시작
        for(int i=0;i<N+1;i++){
            adj[i] = new ArrayList<>();
        }

        for(int i=1;i<=N-1;i++){
            String UV = br.readLine();
            st = new StringTokenizer(UV, " ");
            int U = Integer.parseInt(st.nextToken());
            int V = Integer.parseInt(st.nextToken());

            adj[U].add(V);
            adj[V].add(U);
        }

        //2. child 수 세기
        childCount = new int[N+1];
        Arrays.fill(childCount, 1); //자기자신도 포함
        SubTree(R, -1); //-1 : 부모 없음

        //3. 결과
        sb = new StringBuilder();
        for(int T=1;T<=Q;T++){
            int U = Integer.parseInt(br.readLine());
            sb.append(childCount[U]).append("\n");
        }

        System.out.println(sb);
    }

    public static void SubTree(int child, int parent){

        for(int i=0;i<adj[child].size();i++){
            int next = adj[child].get(i);
            if(next == parent) continue;
            SubTree(next, child);
        }

        if(parent!=-1){//부모가 있다면
            childCount[parent] += childCount[child];
        }
    }
}