第 154 场双周赛——带权树中的最短路径

题目

带权树中的最短路径

题解

我们如果将一个边权更新，那么经过它的路径长度都会发生变化，这很像维护一个差分数组，但是我们如何将它的子树变成一个区间，这里引入了一个DFS时间戳，这个思路很巧的将子树放入了一个区间，我们通过求出每个节点的进出时间来获取每个节点的位置关系。我们这题使用差分树状数组来实现。由于这是边权而不是点权，我们可以将边权值放到两个点中作为孩子的点的身上，因为孩子只有一个父亲，父亲不一定只有一个孩子。这里的时间戳非常巧妙！！！

参考代码

template<typename T>
class FenwickTree {
    vector<T> tree;

public:
    FenwickTree(int n) : tree(n + 1) {}

    void update(int i, T val) {
        for(;i < tree.size(); i += i & -i) tree[i] += val;
    }

    T pre(int i) const {    //const防止数据意外更改
        T res = 0;
        for(; i > 0; i &= i - 1) res += tree[i];
        return res;
    }
};

class Solution {
public:
    vector<int> treeQueries(int n, vector<vector<int>>& edges, vector<vector<int>>& queries) {
        vector<vector<int>> g(n + 1);
        for(auto& e : edges) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }

        //通过时间戳来判断是它的子树
        vector<int> in(n + 1), out(n + 1);
        int clock = 0;
        auto dfs = [&](this auto&& dfs, int x, int fa) -> void {
            in[x] = ++clock;
            for(auto y : g[x]) {
                if(y != fa) {
                    dfs(y, x);
                }
            }
            out[x] = clock;
        };
        dfs(1, -1);

        vector<int> weight(n + 1);  //把边权保存在儿子中
        FenwickTree<int> diff(n);
        auto update = [&](int x, int y, int w) {
            if(in[x] > in[y]) y = x;
            int d = w - weight[y];
            weight[y] = w;
            diff.update(in[y], d);
            diff.update(out[y] + 1, -d);
        };

        for(auto&e : edges) update(e[0], e[1], e[2]);

        vector<int> ans;
        for(auto& q : queries) {
            if(q[0] == 1) update(q[1], q[2], q[3]);
            else ans.push_back(diff.pre(in[q[1]]));
        } 
        return ans;
    }
};