[백준] 15681 트리와 쿼리 풀이 Java

 

15681번: 트리와 쿼리

 

문제

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

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

(*) 메모리 제한 : 128MB

(*) 백준 문제 하단에 Tree에 대한 간단한 설명이 있으니 읽어보면 좋음

 

트리의 정점의 수 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)

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

풀이

DP+Tree 문제

 

처음에는 Tree를 2차원 배열로 구현했는데 메모리초과 발생!

N이 최대 10만까지 가능해서 N*N 배열을 생성하면 .. int 배열일 때 10만x10만 배열을 생성하면

4 bytes * 10만 * 10만 = 약 40GB 

 

ArrayList 배열을 사용하는 것으로 변경!

 

- dfs(R) 를 한번 호출해서 루트부터 시작해 모든 정점에 대한 서브트리의 갯수를 dp 배열에 저장한다.

- U가 주어지면 dp 배열을 사용해서 바로 결과를 리턴한다.

 

import java.util.*;

public class Main {
	static int[] dp;
	static ArrayList<Integer>[] tree;
	static int[] visited;
	static int n;

	public static void main(String[] args) {
		Scanner scan = new Scanner(System.in);

		n = scan.nextInt();
		int r = scan.nextInt();
		int q = scan.nextInt();

		tree = new ArrayList[n + 1];
		dp = new int[n + 1];
		visited = new int[n + 1];

		for (int i = 1; i <= n; i++) {
			tree[i] = new ArrayList<Integer>();
		}

		for (int i = 0; i < n - 1; i++) {
			int x = scan.nextInt();
			int y = scan.nextInt();
			tree[x].add(y);
			tree[y].add(x);
		}

		// Start!
		visited[r] = 1;
		func(r);

		for (int i = 0; i < q; i++) {
			int idx = scan.nextInt();
			System.out.println(dp[idx]);
		}
	}

	public static int func(int idx) {

		int res = 1;

		if (dp[idx] == 0) {
			ArrayList<Integer> list = tree[idx];

			for (Integer val : list) {
				if (visited[val] == 0) {
					visited[val] = 1;
					res += func(val);
				}
			}

			dp[idx] = res;
		}

		return dp[idx];
	}
}