백준 15681 트리와 쿼리 자바

문제는 다음과 같다.
https://www.acmicpc.net/problem/15681

풀이

풀이는 다음과 같다.

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

public class Main {
    static ArrayList<Integer> [] A;
    static int N, R, Q;
    static int[]  subTree;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        StringTokenizer st1 = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st1.nextToken());
        R = Integer.parseInt(st1.nextToken()); //루트 노드 R인 트리구조 (문제에서 트리라고 주어짐)
        Q = Integer.parseInt(st1.nextToken());

        subTree = new int[N+1];
        A = new ArrayList[N+1];
        for(int i = 0 ; i <= N ; i++) {
            A[i] = new ArrayList<Integer>();
        }

        for(int i = 1 ; i <= N-1 ; i++) { //트리이다 -> 간선 N-1개 / 역은 안된다.
            StringTokenizer st2 = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st2.nextToken()); //시작점
            int b = Integer.parseInt(st2.nextToken()); //끝점
            
            A[a].add(b);
            A[b].add(a);
        }

        //해야할것 -> 그래프에서 1과 연결된 노드들을 A[1] 에 넣었는데, 경우의 수는 둘중하나이다.
        //A[1]에 1의 부모노드가 있거나, 없거나
        //그렇다면, 
        //A[1]을 순회하며 걸린 현재노드가 부모노드가 아니라면 서브트리에 속한 노드이므로 숫자를 세주면 된다.
        //visited 없어도 되는 이유도, 트리 특성상 부모노드 -> 자식노드 로 진행이 되는데,
        //어차피 부모 노드만 방문했고 나머지 노드는 미방문 상태이기 때문.
        DFS(-1, R);

        for(int i = 0 ; i < Q ; i++) {
            int a = Integer.parseInt(br.readLine());
            bw.write(String.valueOf(subTree[a]) + "\n");
        }

        bw.flush();
        bw.close();


    }

    static void DFS(int parents, int baby) {
        subTree[baby] = 1; //자기 자신도 서브트리 수에 포함되니까.

        for(int friends : A[baby]) {
            if(friends == parents) continue;
            DFS(baby, friends);
            subTree[baby] += subTree[friends]; //DFS를 빠져나오면서, 누적된다
        }
    }


}

트리구조를 잘 알고있다면, 어렵지 않게 풀 수 있는 문제이다.

근데 나는 트리구조의 특성을 잘 모르고있어서
쉽게 풀 수 있지 않은 문제였다.
트리구조를 여러개 만들건데,
subTree[i] = i값을 루트노드로 하는 트리구조의 크기
라고 하자.

즉,

이러한 트리 구조가 있을때,
index가 3이라면
잡으면 해당 점이 루트 노드가 되고, 3, 1, 2 이렇게 트리에 포함된다.
subTree[3] = 3

마찬가지로
index가 6이면
6이 루트노드가 되고,
6, 7, 8, 9 이렇게 트리에 포함된다.
subTree[6] = 4
루트 노드로 잡는다는 말은, 해당 노드의 상위 노드는 없다고 본다.

왜 이렇게 subTree 배열을 만들었냐?

문제에서 입력값으로 주어진 숫자를 루트 노드로 보고,
해당 숫자를 시작으로 순회하여 노드 개수를 구하면,
subtree를 구성하는 노드 개수를 구할 수 있기 때문이다.

여기까지 이해갔다면, 한걸음 남았다.

static void DFS(int parents, int baby) {
        subTree[baby] = 1; //자기 자신도 서브트리 수에 포함되니까.

        for(int friends : A[baby]) {
            if(friends == parents) continue;
            DFS(baby, friends);
            subTree[baby] += subTree[friends]; //DFS를 빠져나오면서, 누적된다
        }
    }

해당 DFS를 보면, 자식노드(A[baby]) 를 순회하며, 해당 노드와 연결된 노드가 부모노드(patents)가 아닌 이상 다시 DFS를 실행시킨다.
그렇게 자식노드들을 들어가며,
끝에 도달하고 DFS를 빠져나오며,
자식노드의 subTree값을 더해주며 부모노드의 subTree 배열을 채워준다.
(subTree[baby] += subTree[friends];)

다시 이 그림으로 예시를 들면,

그래프의 노드 1과 2는 자식노드가 없으므로 더해지는 값은 없고,
자기 자신을 루트 노드라고 했을때
크기 1인 트리가 되고,
subTree[1] = 1, subTree[2] = 1이 된다.

그 후 DFS를 빠져나오며

자식노드 1과 2의 부모노드 3은
subTree[3] += subTree[1], subTree[3] += subTree[2]를 거치며
subTree[3] = 3 (원래있던1, 더해준값 각각 1)이 된다.

이런식으로 subTree가 채워진다.

마치며 기억하면 좋을 부분

필요한 지식은 트리구조가 제대로 되어있다면
부모노드만 피하면 되니까
visited 배열은 불필요,

그다음 DFS 매서드 에서
개수를 세어주며 빠져나오는 로직정도를 숙지하면 좋을것 같다.

초깃값 subTree[parents] = 1;

그다음 for문으로 baby 순회하며 
DFS(baby, friends);
subTree[baby] += subTree[friends];