第 152 场双周赛——最长特殊路径 II

题目

最长特殊路径 II

题解

在做这题的时候我们先看看他的前身最长特殊路径，这是一道树上滑窗，窗口的大小是由两个相同颜色的节点作为边界，我们可以记录每个颜色最近的节点深度，这样子每次查找时间复杂度就降到了O(1)，我们通过 dis 栈记录从根节点到当前节点的路径权重和。dis 栈的每个元素都表示当前路径上节点的累计权重。而当前选取的路径长度就为 栈顶元素-dis[top_depth]，节点个数就为 len(dis) - top_depth 为了方便代码编写，我们可以把输出整成pair形式，然后每次取max，而节点个数只要取个负值就能获得最小的，省的if判断。这样子我们就知道了他的前身是如何完成的。回看这题，与第一版的区别就在于他可以允许有一个颜色相同一次，那么其实我们只需要将代码稍作修改就可以实现，只需要再拿个变量维护另一个颜色的深度即可。

参考代码(前身)

class Solution {
    vector<vector<pair<int, int>>> g;
    vector<int> nums;
    pair<int, int> ans = {-1, 0};
    vector<int> dis = {0};
    unordered_map<int, int> last_depth;
    void dfs(int x, int fa, int top_depth){
        int color = nums[x];
        int old_depth = last_depth[color];
        top_depth = max(top_depth, old_depth);
        ans = max(ans, pair(dis.back() - dis[top_depth], top_depth - (int) dis.size()));
        last_depth[color] = dis.size();
        for(auto& [y, w] : g[x]){
            if(y != fa){
                dis.push_back(dis.back() + w);
                dfs(y, x, top_depth);
                dis.pop_back();
            }
        }
        last_depth[color] = old_depth;
    }
public:
    vector<int> longestSpecialPath(vector<vector<int>>& edges, vector<int>& nums) {
        g.resize(nums.size());
        for(auto& e : edges){
            int x = e[0], y = e[1], w = e[2];
            g[x].emplace_back(y, w);
            g[y].emplace_back(x, w);
        }
        this->nums = nums;
        dfs(0, -1, 0);
        return {ans.first, -ans.second};
    }
};

参考代码(II)

class Solution {
    vector<vector<pair<int, int>>> g;
    vector<int> nums;
    pair<int, int> ans = {-1, 0};
    vector<int> dis = {0};
    unordered_map<int, int> last_depth;
    void dfs(int x, int fa, int top_depth, int last1){
        int color = nums[x];
        int last2 = last_depth[color];
        top_depth = max(top_depth, min(last1, last2));
        ans = max(ans, pair(dis.back() - dis[top_depth], top_depth - (int) dis.size()));
        last_depth[color] = dis.size();
        for(auto& [y, w] : g[x]){
            if(y != fa){
                dis.push_back(dis.back() + w);
                dfs(y, x, top_depth, max(last1, last2));
                dis.pop_back();
            }
        }
        last_depth[color] = last2;
    }
public:
    vector<int> longestSpecialPath(vector<vector<int>>& edges, vector<int>& nums) {
        g.resize(nums.size());
        for(auto& e : edges){
            int x = e[0], y = e[1], w = e[2];
            g[x].emplace_back(y, w);
            g[y].emplace_back(x, w);
        }
        this->nums = nums;
        dfs(0, -1, 0, 0);
        return {ans.first, -ans.second};
    }
};