[백준 / BOJ] 23089 사탕나무

문제

출처 : https://www.acmicpc.net/problem/23089

23089번: 사탕나무

트리구조의 사탕나무에 사탕이 열려있다.

안즈는 기준이 되는 사탕하나를 골라 해당 사탕부터 K이하에 있는 모든 사탕을 먹을려한다. 이때, 최대로 먹을 수 있는 사탕의 갯수를 출력하는 프로그램을 만드는 문제다.

---

풀이

누적 합에 DP를 이용해 푼 문제다.

문제 풀이 자체는 어렵지 않았으나, 구현시 신경써야할점이 많았었다.

 

구현

문제에서 어려웠던것은, K가 1초과일때(K가 1초과인 이유는 K가 1이라면, 자신의 자식만 확인하면되므로 문제되지 않기 때문이다.), K번째 자식까지 선택한 사탕의 갯수를 빠르게 찾기 힘들다는 점 이였다.

이 를 빠르게 찾기 위해서 누적합에 DP를 사용했다.

 

DP배열을 다음과 같이 선언하고, 트리를 바텀업 방식으로 탐색한다고 가정하자.

dp[idx][K] = idx번째 사탕나무에서 1...K까지의 사탕을 모두 먹은 누적합 (0 <= K <= 20)

리프노드에서는 자기자신만 먹을 수 있으므로, dp[leaf][1] , dp[leaf][2] ... dp[leaf][K-1], dp[leaf][K]는 모두 1로 초기화 된다.

 

부모노드에서는 K=?일때, 자신의 ?번째 자식 사탕을  먹을 수 있다. 따라서, 누적합에 의해 다음과 같이 초기화 된다.

K = 1, dp[parent][1] += dp[leaf][0]

K = 2, dp[parent][2] += dp[leaf][1]

.

.

.

K = K, dp[parent][K] += dp[leaf][K-1]

 

위 탐색을 바텀업 방식으로 진행하면, (부모노드가 없는 시작 노드를 제외한)모든 노드들이 자신의 부모 방향으로 K번째 자식사탕을 먹는 경우를 제외하고는 모두 최대치로 초기화 된다.

 

소스코드

	private int setDp(int idx){
		visit[idx] = true;
		orders.add(idx);
		for(int i = 0; i <= K; i++) dp[idx][i] = 1;
		List<Integer> list = graph.get(idx);
		for(int i = 0; i < list.size(); i++){
			int nextIdx = list.get(i);
			if(visit[nextIdx]) continue;
			par[nextIdx] = idx;
			setDp(nextIdx);
			for(int j = 1; j <= K; j++) dp[idx][j] += dp[nextIdx][j-1];
		}
		return dp[idx][K];
	}

이제, 자신의 부모방향으로 K번째 자식사탕을 먹는경우만 더해주면 된다.

이를 위해, par배열을 만들고, 자신의 부모 idx로 초기화 해준다.

par[idx] = 자신의 부모 인덱스;

 

자신의 부모방향으로 K번째 자식사탕을 먹는경우는 다음과 같이 구할수있다.

1. 처음 dp배열을 초기화하는 방식과 마찬가지로 dp[부모인덱스][K-1]의 값을 더해준다.

2. 이때, 1번과정에서 자신의 방향으로 오는 경우가 더해졌으므로, 부모노드에서 자신으로 향하는 경우를 빼준다.

 

단, 이 과정에서 dp배열을 초기화하는 순서와 동일하게 반복을 진행해줘야한다. 부모노드를 업데이트하지않고, 자식 노드를 업데이트 하면 자식노드에 부모노드의 방향이 반영되지 않을 수 있기 때문이다.

 

소스코드

	private void setPar(){
		orders.poll();
		while(!orders.isEmpty()){
			int idx = orders.poll();
			int[] tmp = new int[K+1];
			for(int i = 0; i <= K; i++) tmp[i] = dp[idx][i];
			for(int i = 1; i <= K; i++){
				if(i == 1) dp[idx][i] += 1;
				else dp[idx][i] += dp[par[idx]][i-1] - tmp[i-2];
			}
		}
	}

 

 

이제, dp배열에서 최댓값을 찾아 출력하면 된다.

 

시간복잡도는 

dp 초기화 : O(NK)

부모 경로 더해주는 연산 : O(NK)

O(2NK) 가 된다.

---

 

전체 소스코드

https://github.com/devxb/JJUNalgo/blob/master/20389%20%EC%82%AC%ED%83%95%EB%82%98%EB%AC%B4/Main.java

GitHub - devxb/JJUNalgo: 백준 알고리즘 소스코드🙃