(Java) 백준 15681

Lee Yechan·2025년 5월 1일

알고리즘 문제 풀이

목록 보기

65/66

시간 제한	메모리 제한	제출	정답	맞힌 사람	정답 비율
1 초	128 MB	20650	9838	7438	45.219%

문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

입력

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)

이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.

이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)

입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.

출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

답안

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.StringTokenizer;

public class BOJ15681 {
    public static void main(String[] args) throws IOException {
		// 첫째 줄 입력 받기
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int nodeCount = Integer.parseInt(st.nextToken());
        int rootNode = Integer.parseInt(st.nextToken());
        int queryCount = Integer.parseInt(st.nextToken());
        
        // 간선 정보 입력 받기
        // 간선 정보를 저장하기 위해 Hashtable을 이용했다.
        Hashtable<Integer, ArrayList<Integer>> edges = new Hashtable<>();
        for (int i = 0; i < nodeCount+1; i++) {
            edges.put(i, new ArrayList<>());
        }
        for (int i = 0; i < nodeCount-1; i++) {
            st = new StringTokenizer(br.readLine());
            int startNode = Integer.parseInt(st.nextToken());
            int endNode = Integer.parseInt(st.nextToken());
            // 문제에서 주어진 그래프는 무방향 그래프이다. 양방향으로 구현하였다.
            edges.get(startNode).add(endNode);
            edges.get(endNode).add(startNode);
        }
        
        // 노드를 생성하고, 그 사이를 edge로 이어 트리를 만든다.
        Node[] nodes = new Node[nodeCount+1];
        for (int i = 0; i < nodeCount+1; i++) {
            nodes[i] = new Node(i);
        }
        for (int i = 0; i < nodeCount+1; i++) {
            for (int endNode : edges.get(i)) {
                nodes[i].neighbors.add(nodes[endNode]);
            }
        }

		// DFS
        dfs(rootNode, 0, nodes);

		// 이미 계산된 정보를 그대로 출력한다.
        for (int i = 0; i < queryCount; i++) {
            int query = Integer.parseInt(br.readLine());
            System.out.println(nodes[query].descendantCount);
        }
    }

    static int dfs(int currentNode, int prevNode, Node[] nodes) {
        for (Node neighbor: nodes[currentNode].neighbors) {
		    // 부모 방향으로의 역방향 탐색 막기
            if (neighbor.nodeNo == prevNode) {
                continue;
            }
            // 자신 노드의 서브트리 노드 개수 = children 노드의 서브트리 노드 개수의 합
            nodes[currentNode].descendantCount += dfs(neighbor.nodeNo, currentNode, nodes);
        }
        return nodes[currentNode].descendantCount;
    }

    static class Node {
        public int nodeNo;  // 노드 번호
        public int descendantCount;  // 자신 노드를 포함한, 서브트리의 노드 개수
        ArrayList<Node> neighbors;  // 자신 노드와 이어진 모든 노드. 부모 포함.

        public Node(int nodeNo) {
            this.nodeNo = nodeNo;
            descendantCount = 1;
            neighbors = new ArrayList<>();
        }
    }
}

풀이

DFS를 통해 트리를 탐색하면서, 트리의 아랫 부분에서부터 각 노드를 루트로 하는 서브트리에 속한 노드의 개수를 세어나갔다.

static class Node {
    public int nodeNo;  // 노드 번호
    public int descendantCount;  // 자신 노드를 포함한, 서브트리의 노드 개수
    ArrayList<Node> neighbors;  // 자신 노드와 이어진 모든 노드. 부모 포함.

    public Node(int nodeNo) {
        this.nodeNo = nodeNo;
        descendantCount = 1;
        neighbors = new ArrayList<>();
    }
}

Node는 트리의 정점을 나타내며, 각 노드는 해당 노드를 루트로 하는 서브트리에 속한 노드의 개수 descendantCount 라는 변수를 가진다.

서브트리에는 자기 자신이 포함되므로, 1로 초기화된다.

static int dfs(int currentNode, int prevNode, Node[] nodes) {
    for (Node neighbor: nodes[currentNode].neighbors) {
        // 부모 방향으로의 역방향 탐색 막기
        if (neighbor.nodeNo == prevNode) {
            continue;
        }
        // 자신 노드의 서브트리 노드 개수 = children 노드의 서브트리 노드 개수의 합
        nodes[currentNode].descendantCount += dfs(neighbor.nodeNo, currentNode, nodes);
    }
    return nodes[currentNode].descendantCount;
}