Tree III

class Solution {
public:
    /* 使用邻接表来表示每一个节点的所有子节点。 */
    vector<vector<int>> graph;
    
    /* height用来存储当前节点向下的最高距离和第二高的距离。 */
    vector<pair<int, int>> height;
    
    /* max1和max2分别表示整体的第一长距离和第二长距离。 */
    int max1, max2;

    int tree3(vector<int>& e) {
        // write code here
        int n = e.size() + 1, i, j, k;
        graph.resize(n + 1);
        
        /* 初始化所有数据。 */
        for(i = 0; i < n - 1; i++) graph[e[i]].push_back(i + 2);
        height.resize(n + 1, {0, 0});
        max1 = 0, max2 = 0;
        
        /* 从root开始进行dfs。 */
        dfs(1);
        return max2;
    }

    void dfs(int start) {
        /* 没有子节点就记录为(1, 0) */
        if(graph[start].size() == 0) height[start] = make_pair(1, 0);
        else {
            /* h1和h2表示当前节点向下的最高距离和第二高的距离。 */
            int h1 = 0, h2 = 0, i, j;
            /* 有子节点，对所有子节点进行dfs。 */
            for(i = 0; i < graph[start].size(); i++) {
                dfs(graph[start][i]);
                if(height[graph[start][i]].first + 1 > h1) {
                    h2 = h1;
                    h1 = height[graph[start][i]].first + 1;
                }
                else if(height[graph[start][i]].first + 1 > h2) h2 = height[graph[start][i]].first + 1;
                if(height[graph[start][i]].second + 1 > h2) h2 = height[graph[start][i]].second + 1;
            }
            /* 如果只有一个子节点，那么就直接更新max1和max2。 */
            if(graph[start].size() == 1) {
                int t1 = height[graph[start][0]].first;
                int t2 = height[graph[start][0]].second;
                if(t1 > max1) {
                    max2 = max1;
                    max1 = t1;
                }
                else if(t1 > max2) max2 = t1;
                if(t2 > max2) max2 = t2;
            }
            /* 如果有两个及以上的子节点，就两层遍历更新max1和max2。 */
            for(i = 0; i < graph[start].size() - 1; i++) {
                for(j = i + 1; j < graph[start].size(); j++) {
                    int tmp1 = height[graph[start][i]].first + height[graph[start][j]].first;
                    int tmp2 = max(height[graph[start][i]].first + height[graph[start][j]].second, height[graph[start][j]].first + height[graph[start][i]].second);
                    if(tmp1 > max1) {
                        max2 = max1;
                        max1 = tmp1;
                    }
                    else if(tmp1 > max2) max2 = tmp1;
                    if(tmp2 > max2) max2 = tmp2;
                }
            }
            /* 更新当前节点的height。 */
            height[start] = make_pair(h1, max(h1 - 1, h2));
        }
    }
};