알고리즘

자바 - 백준 15681 / 트리와 쿼리

kdozlo 2024. 9. 6. 18:21

백준 15681 / 트리와 쿼리
골드5

구현 방법

구현 방법은 루트 노드에서 시작해서 dfs를 이용해 자식노드들을 모두 방문하면서 노드 수를 지속적으로 더하면서 구현했습니다.
이렇게 하면 모든 노드 방문이 끝남과 동시에 각 노드를 루트로 하는 서브 트리의 정점 수를 알 수 있게 됩니다.
따라서 현재 문제의 노드 수(N), 쿼리 수(Q)의 범위가(2 ≤ N ≤ 10^5, 1 ≤ Q ≤ 10^5) 커도 O(N + Q)의 시간복잡도로 해결 가능합니다.
구현 방법은 쉬운편인데 이 아이디어를 생각하는게 꽤 어려웠습니다. 아마 트리 특징에 익숙하지 않아서 그런것 같습니다.

트리 특징

  1. 사이클이 존재 하지 않습니다.(임의의 두 정점 U와 V에 대해, U에서 V로 가는 최단경로는 유일)
  2. 간선수 = 노드수 - 1
  3. 아무 정점이나 잡고 부모와의 연결을 끊었을 때, 해당 정점과 그 자식들, 그 자식들의 자식들… 로 이루어진 부분그래프는 트리가 된다.

의문점

  1. 문제 조건에서 "입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다."를 준 이유는 무엇인가?
    -> 가장 큰 이유는 문제에서 요구하는 서브트리의 정점수를 구할 때 필요하다고 생각합니다.
    다시 말해, 서브트리를 정의하기 위해선 현재 그래프가 트리임이 명확해야 하기 때문입니다.
    -> 두번째는 간선의 수 정보를 따로 주지 않아도 됩니다. 이는 트리의 특징으로 당연하다고 볼 수 있습니다.
  2. 만약 "입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다."라는 조건이 주어지지 않는다면 어떻게 될까?
    -> 서브트리를 구하기 전에 트리인지 판단을 하고 트리가 아닌 경우 트리 형태로 바꾸는 작업이 선행되어야 한다고 생각했습니다.
    이를 판단하고 트리로 바꾸는 작업은 신장 트리 알고리즘(크루스칼, 프림)로 가능합니다.

 

코드

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

/**
 * 백준 15681
 * 트리와 쿼리
 * 골드5
 * https://www.acmicpc.net/problem/15681
 */
public class Boj15681 {

    static int N; // 트리 정점의 수
    static int R; // 루트 노드 번호
    static int Q; // 쿼리 수
    static List<Integer>[] l; // 노드에 연결된 노드 정보
    static boolean[] visited;
    static int[] countNode; // 서브트리 정점 수

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        StringBuilder sb = new StringBuilder();

        N = Integer.parseInt(st.nextToken());
        R = Integer.parseInt(st.nextToken());
        Q = Integer.parseInt(st.nextToken());

        l = new ArrayList[N + 1];
        for(int i = 1; i <= N; i++) {
            l[i] = new ArrayList<>();
        }

        for(int i = 0; i < 7; i++) {
            st = new StringTokenizer(br.readLine());

            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            l[a].add(b);
            l[b].add(a);
        }

        countNode = new int[N + 1];
        visited = new boolean[N + 1];
        countSubtreeNodes(R);

        for(int i = 0; i < Q; i++) {
            int cur = Integer.parseInt(br.readLine());
            sb.append(countNode[cur]).append("\n");
        }

        System.out.println(sb);
    }

    public static void countSubtreeNodes(int cur) {
        countNode[cur] = 1;
        visited[cur] = true;
        for(int i : l[cur]) {
            if(visited[i])
                continue;

            countSubtreeNodes(i);
            countNode[cur] += countNode[i];
        }
    }
}

느낀 점

자료구조의 지식이 필요한 알고리즘 문제를 오랫만에 풀어봐서 좀 당황했습니다. 역시 cs 지식이 중요하다는 것을 다시 한번 느꼈습니다.