Baekjoon Online Judge

algorithm practice

- 단계별 문제풀기

28. 트리

대표적인 그래프 종류 중 하나인 트리를 다뤄 봅시다.

Java / Python

3. 트리의 지름

1967번

가중치가 있는 트리의 지름을 구하는 문제

이번 문제는 입력으로 루트가 있는 트리를 가중치가 있는 간선들로 줄 때, 트리의 지름을 구해서 출력하는 프로그램을 작성하는 문제이다.
루트기준으로 dfs를 진행할 때 가장 큰 가중치를 가진 노드를 구합니다.
그 노드를 기준으로 dfs를 또 구합니다.

• Java

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

public class Main {

	static class Node {
		int node, dist;

		public Node(int node, int dist) {
			this.node = node;
			this.dist = dist;
		}
	}

	static ArrayList<Node>[] list;
	static boolean[] visit;
	static int max = 0;
	static int N;
	static int max_idx = 0;

	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
		StringTokenizer st;

		N = Integer.parseInt(br.readLine());

		list = new ArrayList[N + 1];
		for (int i = 0; i <= N; i++)
			list[i] = new ArrayList<>();

		for (int i = 0; i < N - 1; i++) {
			st = new StringTokenizer(br.readLine());

			int parent = Integer.parseInt(st.nextToken());
			int child = Integer.parseInt(st.nextToken());
			int weight = Integer.parseInt(st.nextToken());

			list[parent].add(new Node(child, weight));
			list[child].add(new Node(parent, weight));
		}

		visit = new boolean[N + 1];
		visit[1] = true;
		dfs(1, 0);

		visit = new boolean[N + 1];
		visit[max_idx] = true;
		dfs(max_idx, 0);
		bw.write(max + "\n");

		bw.flush();
		br.close();
		bw.close();
	}

	public static void dfs(int idx, int dist) {
		if (max < dist) {
			max = dist;
			max_idx = idx;
		}

		for (Node n : list[idx]) {
			if (!visit[n.node]) {
				visit[n.node] = true;
				dfs(n.node, dist + n.dist);
			}
		}
	}
}

• Python

from collections import deque
import sys

input = sys.stdin.readline

def bfs(x, mode):
    que = deque()
    que.append(x)
    arr = [-1 for _ in range(N)]
    arr[x] = 0
    while que:
        x = que.popleft()
        for w, nx in tree[x]:
            if arr[nx] == -1:
                arr[nx] = arr[x] + w
                que.append(nx)
    if mode == 1:
        return arr.index(max(arr))
    else:
        return max(arr)

N = int(input())
tree = [[] for _ in range(N)]

for i in range(N-1):
    x, y, w = map(int, input().split())
    tree[x-1].append([w, y-1])
    tree[y-1].append([w, x-1])
print(bfs(bfs(0, 1), 2))