(백준) 트리 수정 _ 12912

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

 

🐢 설명

트리가 있을 때 하나의 간선만을 제거해서 새롭게 연결했을 때 다시 가질 수 있는 최대 지름을 구하면 되는 문제이다. 이때 새로만든 그래프는 트리로서 특성이 유지가 되어야 한다.

 

그럼 생각해보자. 트리에서 어떤 간선을 제거했다면 어떻게 될까?

 

바로 트리 두개가 생기게 된다. 그렇다면 문제는 단순해진다.

N개의 노드로 이뤄진 트리에서 N-1개의 간선들에 대해서 x간선을 지우고, 트리 두개의 지름을 구한 후, 해당 간선의 길이를 더했을때가 해당 간선을 지웠을 때 최대로 가질 수 있는 지름이 된다.

어짜피 트리 + 트리 이므로 어느 노드에 연결할지는 중요하지 않고 간선을 지웠을 때 각 트리가 가지는 최대 지름만 잘 구해주면 된다.

 

문제의 3번째 예제를 통해 봐보자

이런식으로 5-7 간선을 지웠을 때 2410으로 최대 지름을 갖게 된다.

 

그러면 N-1개의 간선 개수 * ( 1번 트리의 지름 구하기 + 2번 트리의 지름 구하기 + 제거된 간선 길이 ) 로 최대 지름을 구할 수 있다.

주의 할 점은 양방향으로 연결했을 때 특정 노드의 간선의 개수가 2개인데 제거되는 간선의 노드라면 해당 노드가 leaf 노드로서 작동해야 된다는 점이다. 이 부분에 대해 주의토록 하자.

또한, 노드번호가 0번부터 시작하므로 전역변수로 노드등을 설정할 때 기본값으로 0이 아닌 음수 값 등으로 설정하는 것을 잊지말자.

 

핵심 로직 부분은 아래 형태이다. 지름을 구하는 방법은 dfs를 두번 돌려서 구하는 일반적인 방식으로 구하면 된다.

for (int edge = 0; edge < N - 1; edge++) {
    Edge curEdge = edges[edge];

    long firstD = treeDiameter(curEdge.from, curEdge.to);
    long secondD = treeDiameter(curEdge.to, curEdge.from);

    MAX_DIAMETER = Math.max(firstD + secondD + curEdge.dist, MAX_DIAMETER);
}

 

🐢코드

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

public class Main {
	static int N;
	static ArrayList<Edge>[] adj;
	static boolean[] visited;
	static Edge[] edges;
	static long[] dp;
	static long MAX_DIAMETER;
	static long MAX_LEN;
	static int LAST;
	static int END;

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

		N = Integer.parseInt(br.readLine());
		adj = new ArrayList[N + 1];
		for (int i = 0; i < N + 1; i++) {
			adj[i] = new ArrayList<>();
		}
		edges = new Edge[N + 1];
		visited = new boolean[N + 1];
		dp = new long[N + 1];

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

			adj[from].add(new Edge(i, from, to, dist));
			adj[to].add(new Edge(i, to, from, dist));
			edges[i] = new Edge(i, from, to, dist);
		}

		for (int edge = 0; edge < N - 1; edge++) {
			Edge curEdge = edges[edge];

			long firstD = treeDiameter(curEdge.from, curEdge.to);
			long secondD = treeDiameter(curEdge.to, curEdge.from);

			MAX_DIAMETER = Math.max(firstD + secondD + curEdge.dist, MAX_DIAMETER);
		}

		System.out.print(MAX_DIAMETER);
	}

	private static long treeDiameter(int from, int to) {
		LAST = -1;
		END = -1;

		// 시작점이 leaf노드가 되는 경우
		if (adj[from].size() == 2) END = from;

		MAX_LEN = 0;
		Arrays.fill(dp, 0);
		Arrays.fill(visited, false);
		visited[to] = true;
		dfs(from);

		// 자신이 leaf노드였다면 LAST를 갱신할 수 없으므로 자기자신으로 설정한다
		if (LAST == -1) LAST = from;
		MAX_LEN = 0;
		Arrays.fill(dp, 0);
		Arrays.fill(visited, false);
		visited[to] = true;
		dfs(LAST);

		return MAX_LEN;
	}

	private static void dfs(int node) {
		visited[node] = true;

		// leaf 또는 제거된 간선으로 인한 leaf노드인 경우 갱신
		if ((adj[node].size() == 1 || node == END) && dp[node] > MAX_LEN) {
			MAX_LEN = Math.max(MAX_LEN, dp[node]);
			LAST = node;
			return;
		}

		for (Edge child : adj[node]) {
			if (visited[child.to]) continue;

			dp[child.to] += dp[node] + child.dist;

			dfs(child.to);
		}
	}

	private static class Edge{
		int number;
		int from;
		int to;
		int dist;

		public Edge(int number, int from, int to, int dist) {
			this.number = number;
			this.from = from;
			this.to = to;
			this.dist = dist;
		}
	}
}

🐢 마무리