题解 | #多叉树的直径# [P2-T1]

核心就是树的最大直径(globalMax)一定是以某一个node为root最长的两个path-to-leaf.
就是普通dfs的同时算路径长度。

时间: O(n), DFS一次
空间: O(n)

import java.util.*;

/**
 * public class Interval {
 *   int start;
 *   int end;
 * }
 */

public class Solution {
  // node_id -> { {nei1, cost1}, {nei2, cost2}, ... }
  Map<Integer, List<int[]>> graph = new HashMap<>();
  int globalMax = 0;
  
  public int solve (int n, Interval[] Tree_edge, int[] Edge_value) {
    // create each node
    for (int i = 0; i < n; i++) {
      graph.put(i, new ArrayList<int[]>());
    }
    // construct graph
    for (int i = 0; i < n-1; i++) {
      // treat tree edge as bi-directional
      graph.get(Tree_edge[i].start)
        .add(new int[]{Tree_edge[i].end, Edge_value[i]});
      graph.get(Tree_edge[i].end)
        .add(new int[]{Tree_edge[i].start, Edge_value[i]});
    }
    // 因为edge是bi-directional, 随便从哪个node开始dfs都一样。这里用了node-0.
    // -1 作为parent
    dfs(0, -1);
    return globalMax;
  }
  
  // returns the max cost path-to-leaf from this root.
  int dfs(int root, int parent) {
    int maxCost = 0;
    int maxCost2 = 0;
    
    for (int[] neiCost : graph.get(root)) {
      // NOTE: BaseCase (i.e. leaf) has only 1 nei, which is the parent
      //       thus leaf will return maxCost = 0 directly.
      if (neiCost[0] == parent) continue;   // don't revisit parent
      
      // recursively finds the max cost path to any leaf
      int cost = dfs(neiCost[0], root) + neiCost[1];
      // keep 2 largest cost
      if (cost >= maxCost) {
        maxCost2 = maxCost;
        maxCost = cost;
      } else if (cost >= maxCost2) {
        maxCost2 = cost;
      }
    }
    // check if the 2 largest path is the global max
    globalMax = Math.max(
      globalMax,
      maxCost + maxCost2
    );

    return maxCost;
  }
}