Written by Daeyeol Sung

[ Algorithm ] 백준 트리의 독립집합 (2213번)

이번 문제는 백준 트리의 독립집합 문제이다. 트리를 구성하고 DFS를 사용하여 DP 배열을 초기화 하는 것이 핵심이다.

문제

문제는 여기에 참조한다.

풀이

우선 첫째 줄에 정점의 개수 n이 주어지고 두번째 줄에서 각 정점의 가중치가 주어진다. 세번째 줄부터 마지막 줄까지는 에지 리스트가 주어진다. 트리에서 에지의 개수는 n-1개로 일정하지만 문제에서 에지 개수에 대해 명시하지 않아서 EOF로 입력을 받았다.

// 각 정점의 가중치를 저장
StringTokenizer st = new StringTokenizer(br.readLine());
for(int i = 1; i <= N; i++) {
	weights[i] = Integer.parseInt(st.nextToken());
}

// 트리의 에지를 입력받아 저장
while(true) {
	String line = br.readLine();
	if(line == null || line.equals("")) break;
	
	st = new StringTokenizer(line);
	int a = Integer.parseInt(st.nextToken());
	int b = Integer.parseInt(st.nextToken());
	
	edges.get(a).add(b);
	edges.get(b).add(a);
}

에지 리스트를 입력받았으니 이제 트리를 구성하면 된다. 트리에서는 부모 노드와 자식 노드의 관계가 있기 때문에 에지 리스트 중에서 부모 노드로 삼은 노드를 제외한 나머지 노드들만 자식 노드로 연결해준다.

private static void makeTree(int curr, int parent) {
	for(int child : edges.get(curr)) {
		if(child != parent) {
			childs.get(curr).add(child);
			makeTree(child, curr);
		}
	}
}

이제 트리를 구성했으니 DP 배열을 만들고 값을 업데이트 하면 되는데 모든 정점은 독립집합에 속함, 독립집합에 속하지 않음 이렇게 2가지 상태 중 한가지 상태를 가지므로 DP 배열은 아래와 같이 정의했다.

DP 배열을 업데이트 할 때는 DFS를 사용해야 하는데 자식 노드를 따라가며 탐색해서 값을 업데이트 해야 하기 때문이다.

private static void dfs(int curr) {
	dp[curr][0] = weights[curr]; // 현재 정점이 독립집합에 포함된 경우
	dp[curr][1] = 0; // 현재 정점이 독립집합에 포함되지 않은 경우
	
	for(int child : childs.get(curr)) {
		dfs(child);
		dp[curr][0] += dp[child][1]; // 현재 정점이 독립집합에 포함된 경우 자식 정점은 포함될 수 없음
		dp[curr][1] += Math.max(dp[child][0], dp[child][1]); // 현재 정점이 독립집합에 포함되지 않은 경우 자식 정점은 포함될 수도 안될 수도 있음
	}
}

우선 각 재귀 실행에서 해당 정점이 루트 노드인 경우의 DP 값을 초기화 해주고 백트래킹 하면서 이전 실행의 노드(부모)가 독립집합에 포함된 경우와 포함되지 않은 경우를 나눠서 값을 누적해주면 된다. 필자는 전체 트리의 루트 노드를 1번 노드로 해서 트리를 구성했기 때문에 최대 독립집합의 크기는 dp[1][0]과 dp[1][1]에 저장되어 있다. 값이 2개기 때문에 더 큰 값이 최대 독립집합의 크기가 된다.

makeTree(1, -1);

System.out.println(Math.max(dp[1][0], dp[1][1]));

이제 두번째 출력에서 역추적 결과를 출력해야 하는데 역추적 또한 DFS와 DP 배열에 저장된 값을 기반으로 해서 구하면 된다. dp[n][0]에 n번 정점을 포함했을 경우…의 조건이 있기 때문에 최대 독립집합에 포함된 정점들은 dp[n][0] > dp[n][1] 이어야 한다. 또 생각해봐야 할 문제는 n번 정점이 최대 독립집합에 포함되었다면 n번 정점의 자식들은 최대 독립집합에 포함될 수 없다는 것이다. 전체 트리의 루트 노드부터 탐색을 하면서 dp[n][0] > dp[n][1]인 경우 결과에 포함하면 된다.

private static void trace(int curr, boolean isPossible) {
	// isPossible은 독립집합에 포함이 가능한지 여부
	if(isPossible && dp[curr][0] > dp[curr][1]) {
		result.add(curr);
		for(int child : childs.get(curr)) {
			trace(child, false);
		}
	}
	
	else {
		for(int child : childs.get(curr)) {
			trace(child, true);
		}
	}
}

마지막으로 전체 코드이다.

import java.io.*;
import java.util.*;

public class Main {
	static int[] weights;
	static List<List<Integer>> edges = new ArrayList<>();
	static List<List<Integer>> childs = new ArrayList<>();
	static int[][] dp;
	static List<Integer> result = new ArrayList<>();
	
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		
		int N = Integer.parseInt(br.readLine());
		weights = new int[N+1];
		// n번 정점을 루트로 하는 서브트리에서
		// n번 정점이 독립집합에 포함된 경우와 독립집합에 포함되지 않은 경우의
		// 최대 독립 집합의 크기를 저장
		dp = new int[N+1][2];
		
		for(int i = 0; i <= N; i++) {
			edges.add(new ArrayList<>());
			childs.add(new ArrayList<>());
		}
		
		// 각 정점의 가중치를 저장
		StringTokenizer st = new StringTokenizer(br.readLine());
		for(int i = 1; i <= N; i++) {
			weights[i] = Integer.parseInt(st.nextToken());
		}
		
		// 트리의 에지를 입력받아 저장
		while(true) {
			String line = br.readLine();
			if(line == null || line.equals("")) break;
			
			st = new StringTokenizer(line);
			int a = Integer.parseInt(st.nextToken());
			int b = Integer.parseInt(st.nextToken());
			
			edges.get(a).add(b);
			edges.get(b).add(a);
		}
		
		makeTree(1, -1);
		dfs(1);
		trace(1, true);
		Collections.sort(result);
		
		StringBuilder sb = new StringBuilder();
		for(int i = 0; i < result.size(); i++) {
			sb.append(result.get(i)).append(" ");
		}

	
		System.out.println(Math.max(dp[1][0], dp[1][1]));
		System.out.println(sb);
	}
	
	private static void makeTree(int curr, int parent) {
		for(int child : edges.get(curr)) {
			if(child != parent) {
				childs.get(curr).add(child);
				makeTree(child, curr);
			}
		}
	}
	
	private static void dfs(int curr) {
		dp[curr][0] = weights[curr]; // 현재 정점이 독립집합에 포함된 경우
		dp[curr][1] = 0; // 현재 정점이 독립집합에 포함되지 않은 경우
		
		for(int child : childs.get(curr)) {
			dfs(child);
			dp[curr][0] += dp[child][1]; // 현재 정점이 독립집합에 포함된 경우 자식 정점은 포함될 수 없음
			dp[curr][1] += Math.max(dp[child][0], dp[child][1]); // 현재 정점이 독립집합에 포함되지 않은 경우 자식 정점은 포함될 수도 안될 수도 있음
		}
	}
	
	private static void trace(int curr, boolean isPossible) {
		// isPossible은 독립집합에 포함이 가능한지 여부
		if(isPossible && dp[curr][0] > dp[curr][1]) {
			result.add(curr);
			for(int child : childs.get(curr)) {
				trace(child, false);
			}
		}
		
		else {
			for(int child : childs.get(curr)) {
				trace(child, true);
			}
		}
	}
}

이번 문제와 같은 카테고리로 묶여있는 다른 문제들 또한 모두 DFS와 DP 배열을 조합하면 어렵지 않게 풀 수 있기 때문에 풀어보는 것을 추천한다.