[백준] 15681번 : 트리와 쿼리

문제

 

 

접근 방식

정점이 N개이고 루트가 R인 올바른 트리에 대한 간선이 주어졌을 때

정점이 U인 경우에는 몇 개의 정점을 방문할 수 있는지를 알아내면 된다.

 

 

예를 들어 위와 같은 트리가 있을 때를 생각해보면

루트인 5번 정점부터 방문하면 당연히 모든 정점인

9개의 정점을 다 방문할 수 있다.

 

 

하지만 4번 정점부터 방문한다면 위와 같이 4개의 정점만 방문할 것이고

 

 

8번 정점부터 방문한다면 8번 하나만 방문할 수 있다.

 

이제 특정 정점에 포함된 정점의 개수가 몇 개인지를 알아내면 되는데

이는 트리의 성질을 이용하면 쉽게 계산할 수 있다.

 

 

위의 트리를 다시 살펴보면

정점 3의 경우는 자기 자신을 포함한 자식 노드의 정점 수의 합인 3이고

정점 4의 경우도 마찬가지로 자기 자신을 포함한 자식 노드인 3번 노드의

정점 수의 합인 4가 되고, 5번 정점도 마찬가지로 자기 자신 + 4번 + 6번이 된다.

 

즉, 탐색하면서 부모 노드에 자식 노드들을 더해주면 각 노드에 포함된

모든 정점의 개수를 알 수 있고, 쿼리마다 탐색을 하지 않아도

한 번의 탐색으로 처리할 수 있다.

 

풀이

public class Main {

	private static final ArrayList<ArrayList<Integer>> near = new ArrayList<>();
	private static int[] counts;

	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();

		int n = Integer.parseInt(st.nextToken());
		int r = Integer.parseInt(st.nextToken());
		int q = Integer.parseInt(st.nextToken());
		int x, y;

		counts = new int[n + 1];

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

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

			x = Integer.parseInt(st.nextToken());
			y = Integer.parseInt(st.nextToken());

			near.get(x).add(y);
			near.get(y).add(x);
		}

		travel(r);

		while (q-- > 0) {
			sb.append(counts[Integer.parseInt(br.readLine())]).append("\n");
		}

		System.out.println(sb);
	}

	private static void travel(int cur) {
		counts[cur] = 1;

		for (int next : near.get(cur)) {
			if (counts[next] != 0) continue;
			travel(next);
			counts[cur] += counts[next];
		}
	}
}