재귀 트리 DP / Rerooting

# Rerooting: 모든 루트 트리 높이
import sys
sys.setrecursionlimit(300000)
​
N = int(input())
adj = [[] for _ in range(N + 1)]
for _ in range(N - 1):
  u, v = map(int, input().split())
  adj[u].append(v)
  adj[v].append(u)
​
dp_down = [0] * (N + 1)
dp_up = [0] * (N + 1)
ans = [0] * (N + 1)
​
def dfs_down(u, p):
  dp_down[u] = 0
  for v in adj[u]:
      if v == p:
          continue
      dfs_down(v, u)
      dp_down[u] = max(dp_down[u], dp_down[v] + 1)
​
def dfs_up(u, p, parent_val):
  dp_up[u] = parent_val
  ans[u] = max(dp_down[u], dp_up[u])
  
  children = [v for v in adj[u] if v != p]
  k = len(children)
  prefix = [0] * (k + 1)
  suffix = [0] * (k + 1)
  
  for i in range(k):
      prefix[i+1] = max(prefix[i], dp_down[children[i]] + 1)
  for i in range(k - 1, -1, -1):
      suffix[i] = max(suffix[i+1], dp_down[children[i]] + 1)
  
  for i, v in enumerate(children):
      contribution = max(dp_up[u] + 1, prefix[i] + 1, suffix[i+1] + 1)
      dfs_up(v, u, contribution)
​
dfs_down(1, 0)
dfs_up(1, 0, 0)
​
for i in range(1, N + 1):
  print(f"ans[{i}] = {ans[i]}")

// Java Rerooting
import java.util.*;
import java.io.*;
​
public class Main {
  static List<Integer>[] adj;
  static int[] dpDown, dpUp, ans;
  static int N;
  
  static void dfsDown(int u, int p) {
      dpDown[u] = 0;
      for (int v : adj[u]) {
          if (v == p) continue;
          dfsDown(v, u);
          dpDown[u] = Math.max(dpDown[u], dpDown[v] + 1);
      }
  }
  
  static void dfsUp(int u, int p, int parentVal) {
      dpUp[u] = parentVal;
      ans[u] = Math.max(dpDown[u], dpUp[u]);
      
      List<Integer> children = new ArrayList<>();
      for (int v : adj[u]) {
          if (v != p) children.add(v);
      }
      
      int k = children.size();
      int[] prefix = new int[k + 1];
      int[] suffix = new int[k + 1];
      
      for (int i = 0; i < k; i++)
          prefix[i+1] = Math.max(prefix[i], dpDown[children.get(i)] + 1);
      for (int i = k - 1; i >= 0; i--)
          suffix[i] = Math.max(suffix[i+1], dpDown[children.get(i)] + 1);
      
      for (int i = 0; i < k; i++) {
          int v = children.get(i);
          int contribution = Math.max(Math.max(dpUp[u] + 1, prefix[i] + 1), suffix[i+1] + 1);
          dfsUp(v, u, contribution);
      }
  }
  
  public static void main(String[] args) throws IOException {
      BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
      N = Integer.parseInt(br.readLine());
      adj = new ArrayList[N + 1];
      for (int i = 0; i <= N; i++) adj[i] = new ArrayList<>();
      dpDown = new int[N + 1];
      dpUp = new int[N + 1];
      ans = new int[N + 1];
      
      for (int i = 0; i < N - 1; i++) {
          StringTokenizer st = new StringTokenizer(br.readLine());
          int u = Integer.parseInt(st.nextToken());
          int v = Integer.parseInt(st.nextToken());
          adj[u].add(v);
          adj[v].add(u);
      }
      
      dfsDown(1, 0);
      dfsUp(1, 0, 0);
      
      for (int i = 1; i <= N; i++)
          System.out.println("ans[" + i + "] = " + ans[i]);
  }
}