Algorithm

[백준] 2132 - 나무 위의 벌레(Java)

코드파고 2025. 4. 17. 21:14

문제 파악

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

풀이

트리의 지름을 활용해서 푸는 문제이다.
이 문제는 가중치를 트리의 지름에 반영한다는 점이 독특하다.
트리의 지름이면..? DFS를 두 번 수행해야겠다 😊

임의의 정점에서 시작해 가장 많은 열매를 먹을 수 있는 경로를 찾는다.
이 때, 해당 경로의 종점을 저장하자. (1번째 탐색)

저장해 둔 종점에서 다시 트리 순회를 시작해 가장 많은 열매를 먹을 수 있는 경로를 찾는다. (2번째 탐색).

그림으로 보면 다음과 같다. 

첫번째 순회 (파란 원소에서 시작)
두 번째 순회 (노란 종점에서 시작해 종점은 초록)

가장 긴 경로를 찾게 되면 두번째 탐색 경로가 결국 첫 번째 경로를 포함하는 형태가 나올 것이다.

그러므로 두 번째 탐색 시의 시작점, 종점 중 큰 점을 가중치와 함께 반환해 주면 정답이 나온다고 판단했다.

이를 코드로 풀어 내자면 다음과 같다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {
    static int[] fruitStatus;
    static int vertex;
    static List<Integer>[] conn;
    static int maxFruit;
    static int farthestNode;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        vertex = Integer.parseInt(br.readLine());
        fruitStatus = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
        conn = new List[vertex];
        
        // 관계 선언 및 정리
        for (int i = 0; i < conn.length; i++) {
            conn[i] = new ArrayList<>();
        }
        for (int i = 0; i < vertex - 1; i++) {
            String[] input = br.readLine().split(" ");
            int a = Integer.parseInt(input[0]) - 1;
            int b = Integer.parseInt(input[1]) - 1;
            conn[a].add(b);
            conn[b].add(a);
        }

        farthestNode = 0;
        maxFruit = 0;


	// 첫 번째 dfs 수행 -> 가장 긴 경로와 도달하는 끝점을 저장
        dfs(0, -1, fruitStatus[0]);
        int bestStart = farthestNode;
        
        // 두 번째 dfs 수행 -> 끝점에서부터 가장 긴 경로 탐색
        dfs(farthestNode, -1, fruitStatus[farthestNode]);

        System.out.println(String.format("%d %d", maxFruit, Math.min(farthestNode, bestStart) + 1));

    }

    static void dfs(int node, int parent, int acc) {
        if (acc > maxFruit) {
            maxFruit = acc;
            farthestNode = node;
        } else if (acc == maxFruit) {
            farthestNode = Math.min(farthestNode, node);
        }

        for (int next : conn[node]) {
            if (next != parent) {
                dfs(next, node, acc + fruitStatus[next]);
            }
        }
    }
}