BOJ) 트리의 지름 (1967 번)

트리의 지름

1967번: 트리의 지름

 

이진 트리인줄 알고 시간을 조금 낭비했다.

왜 틀렸는지 잘 모르겠어서 게시판 글을 읽다가, 이진 트리가 아닌 케이스를 보고 해결했다.

 

풀이

 

풀이 방법은 간단하다.

각 tree의 연결 정보를 ArrayList에 담아두고 (n이 최대 10,000가 되는데, 배열로는 메모리 초과가 발생할 것임)

dp라고 선언한 배열을 통해 각 노드에서 자식 노드의 최대 weight를 메모이제이션 해주었다.

그리고 답을 찾을 때는 우선순위 큐를 통해, 자식 노드 중 최대 weight를 갖는 두 노드를 연결하여 지름을 구했다.

메모이제이션된 값은 자식 노드가 가지고 있는 최대 weight기 때문에, 각 child로 가는 weight를 더해서 우선순위 큐에 추가해주었다.

우선순위 큐에서 빼낼 때는 자식 노드가 2개 미만으로 갖는 경우가 있기 때문에, 조건문을 통해 걸러주었다. (2개 이상으로 갖는 경우도 for문을 통해 걸러주고 있다.)

 

코드

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Collections;
import java.util.PriorityQueue;
import java.util.StringTokenizer;

public class TreeDiameter_1967 {
    static ArrayList<Node>[] tree;
    static int[] dp;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int n = Integer.parseInt(br.readLine());
        initTree(n);
        for(int i=1; i<n; i++) { // n-1번 간선 정보
            StringTokenizer st = new StringTokenizer(br.readLine());
            int parent = Integer.parseInt(st.nextToken());
            int child = Integer.parseInt(st.nextToken());
            int weight = Integer.parseInt(st.nextToken());
            tree[parent].add(new Node(child, weight));
        }

        br.close();
        solution(n);
    }

    private static void solution(int n) {
        int answer =0;
        initDP(n);
        PriorityQueue<Integer> pq;

        for(int i=1; i<=n; i++) {
            int weight = 0;
            pq = new PriorityQueue<>(Collections.reverseOrder());
            for(Node child : tree[i]) { // 이진트리가 아닐 수 있음
                pq.offer(dp[child.vertex] + child.weight);  // child가 가진 최대 거리 + child와 연결되는 거리
            }
            for(int j=0; j<2; j++) {
                if(!pq.isEmpty()) {
                    weight += pq.poll();
                }
            }
            answer = Math.max(answer, weight);
        }
        System.out.println(answer);
    }

    private static void initDP(int n) {
        dp = new int[n+1];
        final int ZERO =0;
        for(int i=1; i<=n; i++) {
            if(dp[i] == ZERO) { // ZERO가 아닌 경우는 이미 탐색해서 저장됐기 때문에 의미 X
                find(i);
            }
        }
    }

    private static void find(int parent) {
        int max =0;
        for(Node node : tree[parent]) {
            if(dp[node.vertex] == 0) { // 자식 node에 저장된 값이 없다면 다시 진행
                find(node.vertex);
            }
            max = Math.max(dp[node.vertex] + node.weight, max); // 자식의 max 거리와 자식의 거리를 더한 값
        }
        dp[parent] = Math.max(dp[parent],max); // dp 정보 수정
    }

    private static ArrayList[] initTree(int n) {
        tree = new ArrayList[n+1];
        for(int i=1; i<=n; i++) {
            tree[i] = new ArrayList<>();
        }
        return tree;
    }

    private static class Node {
        int vertex, weight;

        public Node(int vertex, int weight) {
            this.vertex = vertex;
            this.weight = weight;
        }
    }
}