문제 바로가기

나의 풀이

이번 풀이는 가장 효율적인 풀이는 아니다.
사실 이 문제에서는 굳이 Node나 Tree를 만들 필요 없이 Graph를 만들어놓고 dfs를 해도 되는 문제다.
다만 나는 Node와 Tree 구현을 좀 더 연습하고자 복잡한 풀이를 쓰게 되었다.

사실 마지막 countAllChildren() 메소드도 처음에는 주어지는 query에 해당하는 노드부터 child를 순회하면서 값을 더했다.

하지만, 이 방식은 시간 초과가 났다.

아무리 생각해도 u와 v가 parent - child 순서라는 보장이 없어 그래프를 우선 만들어 놓기는 해야하므로 위의 메소드를 빠르게 변경하는 수밖에 없다고 생각했다.
(Node 클래스를 굳이 사용하고 싶었다..)

따라서 매번 순회해서 구하는 방식이 아니라 dfs를 사용한 메모이제이션을 활용해서 size 배열에 모든 subset의 노드수를 저장해뒀다.

나의 코드

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

class Main {

	static List<List<Integer>> graph;
	static List<Node> tree;
	static boolean[] visited;
	static int[] size;

	static class Node {
		int number;
		Node parent;
		List<Node> children = new ArrayList<>();

		public Node(int number) {
			this.number = number;
		}

		public void setParent(Node parent) {
			this.parent = parent;
		}

		public void setChild(Node child) {
			this.children.add(child);
		}

		public List<Node> getChildren() {
			return children;
		}

		public int getNumber() {
			return number;
		}
	}

	public static void main(String[] args) throws IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st = new StringTokenizer(br.readLine());
		int n = Integer.parseInt(st.nextToken());
		int r = Integer.parseInt(st.nextToken());
		int q = Integer.parseInt(st.nextToken());

		graph = new ArrayList<>();
		for (int i = 0; i <= n; i++) {
			graph.add(new ArrayList<>());
		}

		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());
			graph.get(u).add(v);
			graph.get(v).add(u);
		}

		tree = new ArrayList<>();
		for (int i = 0; i <= n; i++) {
			tree.add(new Node(i));
		}

		visited = new boolean[n + 1];
		size = new int[n + 1];
		graphToTree(r);
		countAllChildren(tree.get(r));

		for (int i = 0; i < q; i++) {
			int query = Integer.parseInt(br.readLine());
			System.out.println(size[query]);
		}
	}

	private static void graphToTree(int parent) {
		visited[parent] = true;
		for (Integer child : graph.get(parent)) {
			if (visited[child]) {
				continue;
			}
			tree.get(parent).setChild(tree.get(child));
			tree.get(child).setParent(tree.get(parent));
			graphToTree(child);
		}
	}

	private static void countAllChildren(Node parent) {
		size[parent.getNumber()] = 1;
		for (Node child : parent.getChildren()) {
			countAllChildren(child);
			size[parent.getNumber()] += size[child.getNumber()];
		}
	}
}