[백준/Java] 14942 개미

🔗 문제 링크: https://www.acmicpc.net/problem/14942

 

🔍 해결 과정

처음에 이 문제를 접하고 들었던 생각은 dfs를 통해서 각 개미 노드 번호의 부모 노드 번호 정보를 저장한 다음 저장한 부모 노드 정보를 활용하여 개미의 에너지를 사용했을 때 최대로 거슬러 올라갈 수 있는 부모 노드의 번호를 저장하는 방식으로 구현했습니다.

 

과정을 요약하자면 다음과 같습니다.

1️⃣ 개미의 방 구조를 나타내는 이중 연결 리스트 graph를 통해서 방의 구조 형태를 그린다. 이 때 양방향으로 연결한다.
2️⃣ 개미의 에너지 정보 배열 energies, 개미의 부모 정보 배열 parent, 개미의 에너지를 활용하여 가장 가까운 방 번호 저장용 배열 result를 만들어준다. 각각의 배열은 일차원 정수형 배열이다.
3️⃣ dfs를 활용하여 각 개미의 부모를 저장하여 parent의 내용을 완성해준다.
4️⃣ graph, energies, parent 배열들을 전부 활용하여 가장 가까운 방 번호 정보를 각각 result 배열에 저장해준다.

 

코드는 아래의 1️⃣ 번 코드의 내용인데 원래는 시간초과가 났어야 하는 문제인데 테스트케이스가 부족해서 정답처리가 된 것 같습니다.

왜냐하면 n의 최댓값은 10만인데 처음에 작성한 코드는 시간복잡도 O(n^2) 의 구조를 하고 있어 최대 100억 번의 연산을 진행하게 되어 문제에서 제시한 시간 제한 2초(1억번의 연산당 1초)를 훌쩍 넘길 수가 있습니다.

 

그래서 현재의 코드에서 더 나은 방식으로 최적화하는 방법을 고민하게 되었습니다.

코드 1️⃣ 번의 getResult 함수에서 부모 노드로 거슬러 올라가면서 가장 가까운 부모 노드를 계산하는데 불필요한 중복 계산을 최적화하기 위해서 메모이제이션을 적용하였습니다.

 

getResult에서 메모이제이션을 적용한 내용을 설명하자면 다음과 같습니다.

1️⃣ HashMap 자료구조를 활용하여 특정 노드에서 최대로 갈 수 있는 부모 노드의 정보를 미리 저장하는 용도로 활용한다. 특정 노드에서 루트 노드와 가장 가까운 노드 번호를 구하려고 할 때 해당 자료구조에 이미 결과가 저장되어 있다면 그 결과를 활용하는 용도로 사용한다.
2️⃣ 특정 노드 번호에서 부모 노드로의 cost를 계산할 때도 parentCost 배열을 활용하여 특정 노드 번호에서 부모 노드로의 cost를 미리 저장한다.
3️⃣ 비용을 계산할 때 저장된 비용 정보를 활용하여 바로 계산을 함으로써 최적화 한다.

 

---

✏️ 코드

1️⃣ dfs를 통해서 부모 정보를 저장 후 최종 방 번호 계산하는 코드 (시간 복잡도: O(n^2)) (정답 처리된 코드)

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

public class Main {

    private static List<List<int[]>> graph = new ArrayList<>();
    private static int[] energies;
    private static int[] result;
    private static int[] parent;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        int n = Integer.parseInt(br.readLine());
        energies = new int[n + 1]; // 각 노드의 개미당 가지고 있는 에너지
        result = new int[n + 1]; // 결과
        parent = new int[n + 1]; // 노드 번호당 가지고 있는 부모 노드의 번호 정보

        for (int i = 1; i <= n; ++i) {
            parent[i] = i;
        }

        for (int i = 0; i <= n; ++i) {
            graph.add(new ArrayList<>());
        }

        for (int i = 1; i <= n; ++i) {
            energies[i] = Integer.parseInt(br.readLine());
        }

        for (int i = 0; i < n - 1; ++i) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int from = Integer.parseInt(st.nextToken());
            int to = Integer.parseInt(st.nextToken());
            int cost = Integer.parseInt(st.nextToken());

            graph.get(from).add(new int[]{to, cost});
            graph.get(to).add(new int[]{from, cost});
        }

        dfs(1, 0); // 루트 노드, 부모 노드

        for (int i = n; i > 0; --i) {
            result[i] = getResult(i);
        }

        StringBuilder sb = new StringBuilder();
        for (int i = 1; i <= n; ++i) {
            sb.append(result[i]).append("\n");
        }
        bw.write(sb.toString());
        bw.flush();
    }

    private static int getResult(int node) {
        // 현재 노드의 체력을 구하기
        int curEnergy = energies[node];
        while (node != 1) {
            int p = 0;
            // 현재
            for (int i = 0; i < graph.get(node).size(); ++i) {
                int next = graph.get(node).get(i)[0];
                if (next == parent[node]) {
                    p = next;
                    curEnergy -= graph.get(node).get(i)[1];
                    break;
                }
            }

            if (curEnergy == 0) {
                return p;
            }

            if (curEnergy < 0) {
                return node;
            }

            node = p;
        }

        return node;
    }

    private static void dfs(int node, int p) {

        parent[node] = p;

        for (int i = 0; i < graph.get(node).size(); ++i) {
            int next = graph.get(node).get(i)[0];

            if (next == p) {
                continue;
            }

            dfs(next, node); // 다음 노드로 넘어감
        }
    }
}

 

2️⃣ 메모이제이션을 활용하여 부모 노드로의 비용을 미리 저장하여 최적화 (시간 복잡도는 1번과 동일하지만 중복 계산을 제거)

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

public class Main {

    private static List<List<int[]>> graph;
    private static int n;
    private static int[] energies;
    private static int[] parent;
    private static int[] parentCost;
    private static Map<Integer, Integer> parentResultMap = new HashMap<>();

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        n = Integer.parseInt(br.readLine());
        parent = new int[n + 1];
        parentCost = new int[n + 1];
        energies = new int[n + 1];
        graph = new ArrayList<>();
        Arrays.fill(parent, -1);

        // 에너지 입력
        for (int i = 1; i <= n; ++i) {
            energies[i] = Integer.parseInt(br.readLine());
        }

        for (int i = 0; i <= n; ++i) {
            graph.add(new ArrayList<>());
        }

        StringTokenizer st;
        for (int i = 0; i < n - 1; ++i) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            int cost = Integer.parseInt(st.nextToken());

            graph.get(a).add(new int[]{b, cost});
            graph.get(b).add(new int[]{a, cost});
        }

        // dfs를 활용하여 부모 정보를 입력하기
        dfs(1, 0, 0); // 1: 루트 노드, 0: 부모 노드

        int[] result = new int[n + 1];
        // parent, parentCost 정보가 다 저장된 상태
        for (int i = n; i > 0; --i) {
            result[i] = getResult(i);
        }

        StringBuilder sb = new StringBuilder();
        for (int i = 1; i <= n; ++i) {
            sb.append(result[i]).append("\n");
        }

        bw.write(sb.toString());
        bw.flush();
    }

    private static void dfs(int node, int p, int cost) {
        parent[node] = p;
        parentCost[node] = cost; // 현재 노드에서 부모 노드로의 드는 비용

        int size = graph.get(node).size();
        for (int i = 0; i < size; ++i) {
            int[] cur = graph.get(node).get(i);
            int to = cur[0];
            int toCost = cur[1];

            if (to == p) {
                continue;
            }

            dfs(to, node, toCost);
        }
    }

    private static int getResult(int node) {

        if (parentResultMap.containsKey(node)) {
            return parentResultMap.get(node);
        }

        int original = node;
        int curHealth = energies[node];

        while (node != 1) {
            // 여기서 메모이제이션을 활용하기
            int cost = parentCost[node];
            curHealth -= cost;

            if (curHealth == 0) {
                parentResultMap.put(original, parent[node]);
                return parent[node];
            }

            if (curHealth < 0) {
                parentResultMap.put(original, node);
                return node;
            }

            node = parent[node];
        }

        // 루트 노드 1번까지 왔다는 얘기임
        parentResultMap.put(original, 1);
        return 1;
    }
}

 

 

확실히 시간이 단축되는 것을 확인할 수가 있습니다.