Posts 알고리즘 - 트리 DP 문제 풀기(트리에서의 Dynamic Programming)
Post
Cancel

알고리즘 - 트리 DP 문제 풀기(트리에서의 Dynamic Programming)


알고리즘 - 트리 DP 해결하기(트리에서의 Dynamic Programming)


트리 DP

  • 그래프와 관련된 문제를 풀 때에는 다익스트라, dfs 등을 많이 사용하기도 하고, 특정한 조건이 있을 때 2차원배열 DP로 풀기도 했었습니다.
  • 트리(tree)라는 자료구조도 기본적으로 그래프의 일종이며, 서로 다른 두 노드를 잇는 길이 하나밖에 없는 사이클이 없는 그래프입니다. 따라서 그래프처럼 충분히 DP를 적용할 수 있습니다.
  • 이 트리에서 DP를 구현한다는 것은,
    • 예를 들면 특정한 i번째 노드를 루트로 하는 서브 트리에 대해서 i번째 루트 노드를 포함 했을 때와 포함하지 않았을 때 중 조건에 맞는 답을 정의
  • 와 같이 다이나믹 프로그래밍트리 구조에서 생각할 수 있을 때 적용하면 좋습니다.

📄 15681번 - 트리와 쿼리

링크에 들어가 보면 문제 하단부에 자세하게 트리와 관련된 설명을 확인할 수 있습니다.

이 문제는 루트가 R 값으로 특정하게 정해져 있습니다. 입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다고 합니다. 이 때 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 구해야 합니다.

  • 접근법

이 문제는 본격적인 트리DP를 풀기 전, 트리와 DP를 어떻게 구성하면 좋을지 생각할 수 있는 좋은 문제입니다.

결론부터 말하자면 DFS + DP 로 풀면 됩니다. 그런데 트리 구조를 가졌다고 해서 코드상에서 트리 자료구조를 직접 구현할 필요는 없습니다. 트리 DP는 대부분 DFS를 통해 리프 노드까지 향한 다음, 올라오면서 루트의 값을 처리하는 형태를 가지고 있습니다.

이 때, 일반적인 DFS와는 다르게 한 번 방문한 곳은 visit을 복구해 주지 않고 다시는 가지 못하게 한다면 트리를 탐색하는 것과 같은 효과를 볼 수 있습니다.

그림으로 한번 살펴보겠습니다. 위의 문제에서 주어지는 기본 인풋에서는 루트가 5번 노드였고, 간선 정보는 다음과 같았습니다.

1
2
3
4
5
6
7
8
1 3
4 3
5 4
5 6
6 7
2 3
9 6
6 8

위 인풋은 각 줄에 있는 두 노드를 끝점으로 하는 간선이 트리에 속함을, 즉 서로 연결되었다는 뜻이고, 이를 나타내면 다음과 같습니다. (사실 트리처럼 표시하지 않고 선을 이리저리 옮겨져 모양을 흐트러뜨리면 일반 그래프 모양과 다름없습니다.)

여기서 DFS로 5번부터 차례대로 탐색할 것입니다. 그래프 정보를 2차원 벡터로 받았을 때, 각 노드마다 아래와 같이 연결 정보가 담길 것입니다.

1
2
3
4
5
6
7
8
9
node[1] - 
node[2] -
node[3] - [1] - [2]
node[4] - [3]
node[5] - [4] - [6]
node[6] - [5] - [7] - [9] - [8]
node[7] -
node[8] -
node[9] -

탐색 과정을 그림으로 나타내 보았습니다.

5번부터 탐색을 시작하여 파란색은 이미 탐색 및 방문을 끝마친 노드를 표시했고, 이제 6번 노드로 가서 탐색을 이어가려는 상황입니다.

앞에서 이야기했듯이, 특정 노드를 출발점으로 하여 새로운 DFS를 진행한다고 하더라도, 이 문제는 이전에 한 번이라도 방문된 노드는 다시 가지 않도록 하면 트리를 탐색하는 것과 같아집니다.

그림에서 예를 들면, 벡터 상에서는 node[6] 에는 5, 7, 9, 8 번이 연결되어 있지만 6과 연결된 5번은 이미 방문한 적이 있기 때문에 가지 않습니다. 즉, 5번 노드를 트리로 설정하고 탐색을 시작했을 때, 6번 노드를 시작으로 하는 탐색을 진행할 때에는 자연스럽게 부모 노드인 5번을 가지 않고 방문하지 않은 자식 노드들(7, 8, 9)을 가게 됩니다.

그리고 이 문제의 답은 각 노드들을 출발점으로 하여 연결된 곳들을 탐색하고 구한 값을 차례대로 누적하여 dp 테이블에 저장하면 됩니다.

초기에 dp 테이블 값은 모두 1로 초기화합니다. 그림에서의 경우, 5-4-3-1 까지 재귀적으로 들어간 뒤 1과 연결된 다음 곳은 없으므로 돌아오고, dp[3]에는 dp[3] + dp[1] 값이 저장됩니다.

이후 3번은 또다른 연결된 노드인 2번에 간 뒤, 2는 연결된 다음 곳이 없으므로 돌아오고, dp[3]에는 dp[3] + dp[2] 값이 저장됩니다. 그러면 dp[3]에는 최종적으로 3이 저장될 것입니다.

이런 식으로 문제를 풀면 됩니다.

  • 📜 소스코드
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int N, R, Q;
int dp[100001];
bool visit[100001];
vector< vector<int> >node;

int dfs(int cur) {
	// 이미 방문한 곳은 가지 않도록 한다.
	if (visit[cur]) return dp[cur];
	visit[cur] = true;		

	// cur와 인접한 노드를 방문한다.
	for (auto next : node[cur]) {
		// 이미 방문한 곳 X
		if (visit[next]) continue;
		dp[cur] = dp[cur] + dfs(next);
	}
	return dp[cur];
}


int main() {

	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	cin >> N >> R >> Q;
	node.resize(N + 1);

	for (int i = 0; i < N - 1; i++) {
		int u, v;
		cin >> u >> v;
		node[u].push_back(v);
		node[v].push_back(u);
	}

	for (int i = 1; i <= N; i++) 
		dp[i] = 1;
	
	dfs(R);
	vector<int> query;
	
	for (int i = 0; i < Q; i++) {
		int U;
		cin >> U;
		cout << dp[U] << '\n';
	}
	return 0;
}
  • (참고) 트리 구조를 갖는 DP 문제는 아니지만 DFS+DP로 푸는 문제로 백준의 내리막길 문제가 있습니다.

📄 1949번 - 우수 마을

본격적으로 트리 DP를 활용해 보겠습니다.

이 문제에서는 마을을 트리 구조로 이루어져 있다고 명시하고 있습니다.마을 하나를 특정한 노드로 보고, 연결 관계를 트리의 간선으로 생각할 수 있습니다.

이 때 조건은 1. 우수 마을로 선정된 마을 주민 수의 총합이 최대여야 하고, 2. 우수 마을 간에는 서로 인접하면 안 되며, 3. 우수 마을로 선정되지 못한 일반 마을은 적어도 하나의 우수 마을과 인접해야 합니다.

이 조건을 트리 관점에서 보면, 우수 마을 이 부모 노드일 경우, 자식 노드에는 우수 마을이 있으면 안 됩니다.

반면 일반 마을이 부모 노드일 경우, 자식 노드에는 우수 마을이 있어야 합니다.

dp 테이블로 바꿔 보면 다음과 같습니다.

dp[cur][0]은 cur번 마을이 일반 마을이었을 때 cur번 마을을 루트로 하는 서브 트리의 우수 마을들 인구수의 최댓값

dp[cur][1]은 cur번 마을이 우수 마을이었을 때 cur번 마을을 루트로 하는 서브 트리의 우수 마을들 인구수의 최댓값

그런 다음 아무 노드나 루트로 삼고 DFS를 진행합니다. 여기서는 1번부터 탐색을 시작했습니다.

탐색할 때는 다음과 같이 dp 테이블을 갱신해 나가면 됩니다.

1
2
3
4
5
6
// 일반 마을은 자식 마을이 우수거나, 자식 마을이 일반 마을이다.
// 더 큰 쪽을 선택한다. 문제의 3번 조건은 여기서 해결이 된다.
dp[cur][0] = dp[cur][0] + max(dp[next][0], dp[next][1]);

// 현재 마을이 우수 마을이면, 자식은 반드시 일반 마을이어야 한다.
dp[cur][1] = dp[cur][1] + dp[cur][0];

또한 이 문제도, 트리 구조를 직접 구현하지 않으면서도 트리 탐색과 같으려면, 한 번 방문한 마을은 다시는 가지 않도록 이전 문제처럼 visit 배열을 통해 관리하면 됩니다.

최종적으로는 루트로 삼았던 마을의 dp[cur][0] 값과 dp[cur][1] 값 중 더 큰 값이 정답입니다.

  • 📜 소스코드
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
#include <deque>

using namespace std;

int N;
int W[10001];
int dp[10001][2];
bool visit[10001];

vector< vector<int> > v;

void dfs(int cur) {

	// 이미 방문한 곳은 가면 안 된다.
	if (visit[cur]) return;

	// 처음 들어온 cur 노드에 대해 일반 마을과 우수 마을일 때 인구수를 초기화
	// 우수 마을을 구하는 것이므로 인구수는 우수마을일 때만 설정한다.
	visit[cur] = true;
	dp[cur][0] = 0;
	dp[cur][1] = W[cur];

	// cur에 인접한 자식 노드들에 대하여
	for (auto next : v[cur]) {
		if (visit[next]) continue; // 방문 안 한 곳에만 간다.

		dfs(next); // 끝까지 타고 들어간다.
		
		// 일반 마을은 자식 마을이 우수거나, 자식 마을이 일반 마을이다.
		// 더 큰 쪽을 선택한다. 3번 조건은 여기서 해결이 된다.
		dp[cur][0] = dp[cur][0] + max(dp[next][0], dp[next][1]);

		// 현재 마을이 우수 마을이면, 자식은 반드시 일반 마을이어야 한다.
		dp[cur][1] = dp[cur][1] + dp[next][0];
	}

}

int main() {

	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);

	cin >> N;

	for (int i = 1; i <= N; i++) {
		int input;
		cin >> input;
		W[i] = input;
	}

	v.resize(N + 1);

	for (int i = 0; i < N - 1; i++) {
		int a, b;
		cin >> a >> b;
		v[a].push_back(b);
		v[b].push_back(a);
	}

	dfs(1);

	cout << max(dp[1][0], dp[1][1]);

	return 0;
}

또다른 트리 DP 문제도 추후 게시판에 업로드할 예정입니다.

GitHub에서 진행하는 코딩테스트 스터디그룹을 소개합니다!

[백준] 2580번 - 스도쿠