题解 | #多叉树的直径#

/**
 * struct Interval {
 *  int start;
 *  int end;
 *  Interval(int s, int e) : start(start), end(e) {}
 * };
 */
class Solution {
  private:
    vector<vector<pair<int, int>>> graph; // 邻接表表示，pair<节点, 边权>
    vector<bool> visited;
    int maxDist;
    int farthestNode;

    // DFS函数，用于遍历并找出最远距离
    void dfs(int node, int dist) {
        visited[node] = true;
        if (dist > maxDist) {
            maxDist = dist;
            farthestNode = node;
        }

        for (auto& edge : graph[node]) {
            if (!visited[edge.first]) {
                dfs(edge.first, dist + edge.second);
            }
        }
    }
  public:
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     * 树的直径
     * @param n int整型 树的节点个数
     * @param Tree_edge Interval类vector 树的边
     * @param Edge_value int整型vector 边的权值
     * @return int整型
     */
    int solve(int n, vector<Interval>& Tree_edge, vector<int>& Edge_value) {
        // write code here
        // 检查输入是不是有问题
        // 以输入的第二部分为基准，
        // 第一部分的数字大或者小不影响，第二部分有 n 对点对，在一棵树里节点数就是 n + 1。
        // 第三部分多或者少也同样不影响，少了我们补 0，多了就不读入多余那部分就好了。
        n = Tree_edge.size() + 1; // 节点数应为边数+1
        if (Edge_value.size() < Tree_edge.size()) {
            Edge_value.resize(Tree_edge.size(), 0); // 边权不足时补0
        }

        // 处理完输入之后再建图
        graph.resize(n);
        visited.resize(n, false);

        for (int i = 0; i < Tree_edge.size(); i++) {
            int u = Tree_edge[i].start;
            int v = Tree_edge[i].end;
            int weight = Edge_value[i];
            graph[u].push_back({v, weight});
            graph[v].push_back({u, weight});
        }

        // 第一次DFS找到最远的节点
        maxDist = 0;
        dfs(0, 0);

        // 第二次DFS从最远的节点开始
        fill(visited.begin(), visited.end(), false);
        maxDist = 0;
        dfs(farthestNode, 0);

        return maxDist;
    }
};