[백준, BOJ 2058] 원자의 에너지 (java)

https://www.acmicpc.net/problem/2058

메모리: 16,212 KB , 시간: 140 ms

사용 알고리즘: 다이나믹 프로그래밍, 트리에서의 다이나믹 프로그래밍, 트리

---

이 문제는 구현보다 문제 자체를 이해하는 게 더 어려웠다.

 

문제를 간단히 정리하면,

원자의 에너지 상태 = 노드라고 보고

'어떤 에너지 상태(A)에서 다른 에너지 상태(B)로 변할 수 있는 방법은 한 가지뿐'이라는 문장으로

어떤 노드(A)에서 다른 노드(B)로 갈 수 있는 방법이 하나뿐인 트리 문제라는 것을 알 수 있다.

 

어떤 노드(A) 값에서 어떤 양성자를 더하거나 빼서 다른 노드(B) 값이 나온다면 두 노드는 서로 연결되어 있다.

1. A 노드에서 모든 양성자의 값을 더하거나 빼 본다.
2. A 노드에서 양성자 값을 더하거나 뺀 값이 다른 노드 값과 동일하다면 A 노드와 B 노드는 연결되어 있다.
3. B노드로 이동 후, B 노드를 대상으로 1번 작업을 수행한다.

아무 노드나 임의로 루트 노드로 잡고 위의 작업을 수행하면, 트리 구조를 파악할 수 있다.

 

이젠, 파악한 트리 구조를 가지고 컨테이너 안에 집어넣을 최댓값을 찾아야 한다.

DFS를 사용하여 최대값을 구할 것인데,

'어떤 한 개의 원자가 한 개의 양성자를 받아들이거나 내쏘아서 다른 원자와 같은 에너지 상태에 도달할 수 있다면, 이 경우는 실험에서 위험이 발생할 수 있다'라는 문제 규칙 상 서로 연결된 노드는 함께 컨테이너 안에 넣을 수 없다.

따라서 자식 노드가 자신을 포함한 최댓값과 자신을 포함하지 않은 최댓값을 배열로 보내주면,

현재 노드도 자신을 포함한 최댓값과 자신을 포함하지 않은 최댓값을 배열로 리턴하도록 구현한다.

 

자신을 포함한 최댓값을 구하기 위해선, 연결된 노드는 함께 컨테이너 안에 넣을 수 없다는 규칙 때문에 자식 노드는 무조건 포함되지 않아야 한다.

자신을 포함한 최댓값 = 자신의 값 + 자식 노드를 포함하지 않은 최대값

 

자신이 포함되지 않은 최대값을 구하기 위해선, 어차피 자신이 포함되지 않으니 자식 노드가 포함되던 말던 상관 없으니 둘 중 더 큰 수를 넣어준다.

자신을 포함하지 않은 최대값 = Math.max(자식 노드를 포함한 최대값, 자식 노드를 포함하지 않은 최대값)

 

이런 로직으로 코드를 구현하면 된다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.HashSet;
import java.util.Set;
import java.util.StringTokenizer;

/**
 * A -> B로 변할 수 있는 방법은 한 가지 뿐이라는 점에서 트리인 것을 알 수 있다.
 *
 * N개의 에너지 상태가 주어지므로, 노드는 N개
 * 만약 노드 A에서 양성자 중 하나를 더하거나 빼서 노드 B값이 나오면 두 노드는 간선으로 연결된 것(A->B)
 *
 * 간선으로 연결된 두 개의 노드는 함께 컨테이너에 집어넣을 수 없음
 * DP를 통해 현재 노드를 넣거나 뺀 경우 중 더 나은 경우를 구해야함
 */
public class Main {

    static int M;
    static boolean[] state;
    static int[] energy;
    static boolean[] visited;

    public static void main(String[] args) throws Exception{

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken()); // 원자 상태의 개수
        M = Integer.parseInt(st.nextToken()); // 양성자의 개수

        // 원자가 가질 수 있는 에너지 상태
        state = new boolean[1_000_001];
        // 임의로 정한 루트 노드의 값
        int start = 0;
        int n;
        for (int i = 0; i < N; i++) {
            n = Integer.parseInt(br.readLine());

            if(i == 0) start = n;
            state[n] = true;
        }

        // 양성자가 가진 에너지
        energy = new int[M];
        for (int i = 0; i < M; i++) {
            energy[i] = Integer.parseInt(br.readLine());
        }

        // 노드 방문 체크
        visited = new boolean[1_000_001];
        
        int[] result = dfs(start);
        System.out.println(Math.max(result[0], result[1]));
    }

    private static int[] dfs(int now) { // now := 현재 상태

        // a -> b -> c -> b -> a와 같은 순환구조를 막기 위해 방문체크
        visited[now] = true;

        // result[0] := now를 포함했을 경우의 최대값
        // result[1] := now를 포함하지 않았을 경우의 최대값
        int[] result = new int[2];
        result[0] = now;

        int[] tmp;
        for (int i = 0; i < M; i++) {
            // 양성자를 받아들이거나 내쏘아서 만들 수 있는 새로운 원자
            for(int newState : new int[] {now + energy[i], now - energy[i]}) {
                if(newState > 0 && newState < 1_000_001) {
                    if(state[newState] && !visited[newState]) { // 만들 수 있는 상태라면(연결된 노드를 찾았다면)
                        // now의 자식노드인 newState의 최대값부터 구하러감
                        tmp = dfs(newState);
                        // result[0]은 now를 포함했을 경우의 최댓값
                        // now와 연결된 newState를 포함하지 않은 경우만 올 수 있음
                        result[0] += tmp[1];
                        // result[1]은 now를 포함하지 않았을 경우의 최대값
                        // now와 연결된 newState를 포함하건 포함하지 않건 상관없음(둘 중 더 큰 값으로)
                        result[1] += Math.max(tmp[0], tmp[1]);
                    }
                }
            }
        }

        return result;
    }
}