백준 15681 트리와 쿼리 [Java]

solved.ac = 골드 5

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

15681번: 트리와 쿼리

 

---

1. 접근

 

특이하게 힌트가 주어진 문제다.

트리의 정의에 대해 잘 모르시는 분은 힌트의 앞부분부터 읽고 이해한 후 문제에 접근하는 게 맞고

그게 아니라면 5번을 정점으로 트리 그림부터 보면 이해가 편하다.

 

문제는 간결하고 힌트는 이론부터 설명하기에 복잡해질 수 있는데 간단히 생각하기 위해 예제 입력을 하나하나 그림으로 뜯어보자.

 

첫째 줄 = 9(정점의 수), 5(루트로 하는 노드의 번호), 3(쿼리의 수)

두번째 줄 ~ 정점의 수 - 1 = 간선 정보

5가 루트인건 아시쥬...!

마지막 쿼리의 수만큼의 줄 = 서브 트리에 속한 정점의 개수를 알고 싶은 노드 번호 (5, 4, 8)

 

첫 번째 쿼리는 5로 입력 루트와 같기에 입력 정점의 수와 동일.

 

두 번째 쿼리는 4로 4번 노드를 정점으로 하는 서브 트리의 정점의 개수

 

마지막 쿼리는 8로 8번 노드를 정점으로 하는 서브 트리의 정점의 개수

 

 

단, 입력 조건을 본다면 매 쿼리마다 독립된 트리를 만들어 탐색하면 시간 초과로 인해 문제를 풀 수 없으니 한 번 탐색하는 김에 각 정점을 루트로 했을 때의 개수를 카운트해두면?

 

 

매 쿼리마다 확인할 필요 없이 만들어둔 배열의 쿼리 번째 인덱스만 호출하면 그 값을 알 수 있다!

 

---

2. 풀이

 

변수 세팅

static int N, R, Q;
static int[] qArr, cnt;
static ArrayList<Integer>[] tree;
public static void main(String[] args) throws IOException {
    st = new StringTokenizer(br.readLine());
    N = Integer.parseInt(st.nextToken());
    R = Integer.parseInt(st.nextToken());
    Q = Integer.parseInt(st.nextToken());

    cnt = new int[N + 1];

    tree = new ArrayList[N + 1];
    for (int i = 1; i <= N; i++) tree[i] = new ArrayList<>();
    for (int i = 1; i < N; i++) {
        st = new StringTokenizer(br.readLine());
        // 문제에서의 그래프는 가중치 없는 무방향 그래프이다.
        int x = Integer.parseInt(st.nextToken()), y = Integer.parseInt(st.nextToken());
        tree[x].add(y);
        tree[y].add(x);
    }

    qArr = new int[Q];
    for (int i = 0; i < Q; i++) qArr[i] = Integer.parseInt(br.readLine());
    solution();
}

 

탐색 및 정답 출력

static void solution() {
    DFS(R, -1); // 시작점은 임의의 루트 R로 한다.
    for (int x : qArr) sb.append(cnt[x]).append("\n");
    System.out.println(sb.toString());
}

 

탐색

static void DFS(int x, int prevNode) {
    cnt[x] = 1; // 루트 또한 간선의 개수에 포함된다.

    for (int y : tree[x]) {
        if (y == prevNode) continue; // 방문한 노드는 재 방문하지 않는다.
        DFS(y, x);
        cnt[x] += cnt[y]; // 상위 노드에 간선의 개수를 누적한다.
    }
}

 

---

3. 전체 코드