[백준] 15681: 트리와 쿼리 - JAVA[자바]

안녕하세요.

오늘은 [백준] 15681: 트리와 쿼리 문제를 풀며 TLE을 경험했습니다...

이번 글에서는 정답코드와 시간초과 코드의 시간복잡도를 분석하고 왜 최적화가 필요했는지 알아보겠습니다.

 

 

🤔 고민했던 점 & 어려웠던 부분

처음 접근 방법은 다음과 같았습니다.

1. 문제에서 정해진 루트(R)을 기준으로 DFS 탐색하여 부모 정보를 기록한다.

2. 쿼리마다 DFS 탐색하여 서브트리의 크기를 계산한다.

// hyeblee
import java.util.*;
import java.io.*;


public class Main {
    
    // 정점 n, 루트 r, 쿼리 수 q
    // 간선 정보 (n-1개) -> 트리니까 단방향으로 연결해야함.
    // 쿼리의 루트노트 q개

    public static List<Integer>[] adjList;
    public static boolean[] visited;
    public static int[] parent;
    public static int n, r, q;
    public static int result = 0;

    // start를 루트로 하는 서브트리의 정점 개수를 구한다.
    public static void dfs(int start) {
        result++;
        visited[start] = true;

        for(int next: adjList[start]) {
            // 인접한 노드가 부모이거나 방문한 적 있으면 탐색하지 않는다.
            if (next == parent[start] || visited[next]) {
                continue;
            }
            dfs(next);
        }
    }

    // r을 루트로 하는 트리의 정점별 부모를 기록한다.
    public static void writeParent(int cur) {
        visited[cur] = true;
        for(int next: adjList[cur]) {
            if (visited[next]) {
                continue;
            }
            parent[next] = cur;
            writeParent(next);
        }
    }


    

    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());
        
        adjList = new ArrayList[n+1];
        for(int i=0;i<n+1;i++) {
            adjList[i] = new ArrayList<>();
        }
        visited = new boolean[n+1];
        parent = new int[n+1];
       

        // 간선 연결하기
        for(int i=0;i<n-1;i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());
            adjList[u].add(v);
            adjList[v].add(u);
        }

        // r을 루트로 하는 트리에서, 각 노드의 정점 기록하기
        parent[r] = r;
        writeParent(r);

        // 쿼리 실행하며 결과 출력하기
        for(int i=0;i<q;i++) {
            int u = Integer.parseInt(br.readLine());

            result = 0;
            Arrays.fill(visited, false);
            dfs(u);
            sb.append(result + "\n");
        }

        System.out.println(sb);
    
    }
}

💥 왜 시간 초과가 발생했을까?

위 코드에서는 쿼리마다 다음과 같은 과정을 반복하고 있습니다.

dfs(u)를 호출하여, 노드 u를 루트로 하는 서브트리의 정점 개수를 매번 새로 계산한다.

 

문제는 이 dfs(u)의 시간복잡도가 O(N) 이라는 점입니다.
트리에서 DFS는 모든 노드를 한 번씩 방문하기 때문이죠.

그런데 이 DFS가 쿼리의 개수 Q만큼 반복되므로, 전체 시간복잡도는 다음과 같습니다:

O(N × Q)

 

⏱️실제 입력 범위로 시간복잡도 분석

• 정점의 수 N ≤ 10⁵
• 쿼리의 수 Q ≤ 10⁵

따라서

O(N × Q) = O(10⁵ × 10⁵) = O(10¹⁰)

 

이는 약 100억 번의 연산으로,
보통 1초에 1억 번 연산이 가능하다고 가정했을 때
약 100초가 걸립니다.

즉, 문제의 시간 제한(1초)을 한참 초과하게 됩니다. 🚨

 

✅ 결론

따라서 매 쿼리마다 DFS를 반복하는 방식은 비효율적입니다.
한 번의 DFS로 모든 노드의 서브트리 크기를 미리 계산해두고,
이후 쿼리마다 O(1) 로 결과를 조회하도록 최적화해야 합니다. (메모이제이션 사용)

 

 

 

📌 올바른 접근

1. 문제에서 정해진 루트(R)을 기준으로 DFS 탐색하여 부모 정보를 기록한다.

2. 한 번의 DFS 탐색(O(n))으로 모든 노드의 서브트리 크기를 계산하고 배열에 저장한다.

3. 쿼리마다 O(1)로 배열을 조회하여 서브트리의 크기를 출력한다. 

 

 

⏱️실제 입력 범위로 시간복잡도 분석

• 부모 노드를 기록하는 DFS: O(N)
• 서브트리 크기를 계산하는 DFS: O(N)
• Q개의 쿼리 처리: O(Q)

따라서 전체 시간 복잡도는

O(N+N+Q)=O(2N+Q)

 

최대 입력 조건인 N=10^5, Q=10^5 를 대입하면,

O(2×10^5+10^5) = O(3×10^5)

 

➡️ 약 30만 번의 연산으로,
일반적으로 1억(10⁸) 연산 ≒ 1초 기준을 고려하면 충분히 1초 이내에 수행됩니다.

 

 

 

 

💻 올바른 풀이

// hyeblee
import java.util.*;
import java.io.*;


public class Main {
    
    // 정점 n, 루트 r, 쿼리 수 q
    // 간선 정보 (n-1개) -> 트리니까 단방향으로 연결해야함.
    // 쿼리의 루트노트 q개

    public static List<Integer>[] adjList;
    public static boolean[] visited;
    public static int[] parent;
    public static int[] subtreeCnt;
    public static int n, r, q;

    // r이 루트인 트리에서 모든 서브트리의 크기를 구한다.
    public static int dfs(int start) {
        int cnt = 1; // 자기 자신
        visited[start] = true;

        for(int next: adjList[start]) {
            // 인접한 노드가 부모이거나 방문한 적 있으면 탐색하지 않는다.
            if (next == parent[start] || visited[next]) {
                continue;
            }
            cnt += dfs(next);
        }
        subtreeCnt[start] = cnt;
        return cnt;
    }

    // r을 루트로 하는 트리의 노드별 부모를 기록한다.
    public static void writeParent(int cur) {
        visited[cur] = true;
        for(int next: adjList[cur]) {
            if (visited[next]) {
                continue;
            }
            parent[next] = cur;
            writeParent(next);
        }
    }


    

    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());
        
        adjList = new ArrayList[n+1];
        for(int i=0;i<n+1;i++) {
            adjList[i] = new ArrayList<>();
        }
        visited = new boolean[n+1];
        parent = new int[n+1];
        subtreeCnt = new int[n+1];
       

        // 간선 연결하기
        for(int i=0;i<n-1;i++) {
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken());
            int v = Integer.parseInt(st.nextToken());
            adjList[u].add(v);
            adjList[v].add(u);
        }

        // r을 루트로 하는 트리에서, 각 노드의 부모 기록하기
        parent[r] = r;
        writeParent(r);
        
        // 한 번의 dfs로 모든 서브트리의 크기 구하기
        Arrays.fill(visited, false);
        dfs(r);

        // 쿼리 실행하며 결과 출력하기
        for(int i=0;i<q;i++) {
            int u = Integer.parseInt(br.readLine());
            int result = subtreeCnt[u];
            sb.append(result + "\n");
        }

        System.out.println(sb);
    
    }
}