3515. 带权树中的最短路径

Hard

题目描述

给定一个整数 n 和一个以节点 1 为根的无向加权树，该树有 n 个节点，编号从 1 到 n。树用长度为 n - 1 的二维数组 edges 表示，其中 edges[i] = [ui, vi, wi] 表示节点 ui 和 vi 之间有一条权重为 wi 的无向边。

给定一个长度为 q 的二维整数数组 queries，其中每个 queries[i] 是以下两种之一：

[1, u, v, w'] - 将节点 u 和 v 之间的边的权重更新为 w'，其中 (u, v) 保证是 edges 中存在的边。
[2, x] - 计算从根节点 1 到节点 x 的最短路径距离。

返回一个整数数组 answer，其中 answer[i] 是第 i 个类型 [2, x] 查询中从节点 1 到 x 的最短路径距离。

示例 1:

输入: n = 2, edges = [[1,2,7]], queries = [[2,2],[1,1,2,4],[2,2]]
输出: [7,4]

示例 2:

输入: n = 3, edges = [[1,2,2],[1,3,4]], queries = [[2,1],[2,3],[1,1,3,7],[2,2],[2,3]]
输出: [0,4,2,7]

提示:

1 <= n <= 10^5
edges.length == n - 1
1 <= wi <= 10^4
1 <= queries.length <= 10^5
边的更新和路径查询操作混合进行

解题思路

这是一道结合树结构、欧拉序和线段树的复杂题目。关键思路是利用欧拉序将树展平成数组，然后用线段树高效处理更新和查询。

核心思想：

欧拉序展平：通过DFS遍历树，为每个节点记录在欧拉序中的进入和退出位置，这样每个节点的子树对应欧拉序中的一个连续区间。
距离维护：维护每个节点到根节点1的距离。当某条边权重改变时，该边下方的所有节点到根的距离都需要调整。
线段树优化：使用线段树支持区间更新和单点查询。当边(u,v)权重从old_w变为new_w时，需要对较深节点的子树进行区间更新，增加new_w - old_w的距离差值。

算法流程：

构建邻接表表示树
DFS获取欧拉序，记录每个节点的进入/退出时间和初始深度
构建线段树，支持区间增加和单点查询
处理查询：类型1更新边权并调整子树距离，类型2查询节点距离

时间复杂度： 每个操作O(log n)，总体O((n+q)log n)

代码实现

class Solution {
public:
    vector<vector<pair<int, int>>> adj;
    vector<int> tin, tout, depth;
    vector<long long> tree;
    vector<long long> lazy;
    int timer;
    int n;
    
    void dfs(int v, int parent, int d) {
        tin[v] = timer++;
        depth[v] = d;
        for (auto &edge : adj[v]) {
            int u = edge.first;
            if (u != parent) {
                dfs(u, v, d + edge.second);
            }
        }
        tout[v] = timer - 1;
    }
    
    void push(int v, int tl, int tr) {
        if (lazy[v] != 0) {
            tree[v] += lazy[v];
            if (tl != tr) {
                lazy[v*2] += lazy[v];
                lazy[v*2+1] += lazy[v];
            }
            lazy[v] = 0;
        }
    }
    
    void update(int v, int tl, int tr, int l, int r, long long val) {
        if (l > r) return;
        if (l == tl && r == tr) {
            lazy[v] += val;
            push(v, tl, tr);
            return;
        }
        push(v, tl, tr);
        int tm = (tl + tr) / 2;
        update(v*2, tl, tm, l, min(r, tm), val);
        update(v*2+1, tm+1, tr, max(l, tm+1), r, val);
    }
    
    long long query(int v, int tl, int tr, int pos) {
        if (tl == tr) {
            push(v, tl, tr);
            return tree[v];
        }
        push(v, tl, tr);
        int tm = (tl + tr) / 2;
        if (pos <= tm)
            return query(v*2, tl, tm, pos);
        else
            return query(v*2+1, tm+1, tr, pos);
    }
    
    vector<int> treeQueries(int n, vector<vector<int>>& edges, vector<vector<int>>& queries) {
        this->n = n;
        adj.resize(n + 1);
        tin.resize(n + 1);
        tout.resize(n + 1);
        depth.resize(n + 1);
        tree.resize(4 * n);
        lazy.resize(4 * n);
        timer = 0;
        
        map<pair<int, int>, int> edgeWeight;
        
        for (auto &edge : edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            adj[u].push_back({v, w});
            adj[v].push_back({u, w});
            edgeWeight[{min(u, v), max(u, v)}] = w;
        }
        
        dfs(1, -1, 0);
        
        vector<int> result;
        
        for (auto &query : queries) {
            if (query[0] == 1) {
                int u = query[1], v = query[2], newW = query[3];
                int oldW = edgeWeight[{min(u, v), max(u, v)}];
                edgeWeight[{min(u, v), max(u, v)}] = newW;
                
                // Find which node is deeper
                int deeper = (depth[u] > depth[v]) ? u : v;
                
                // Update the subtree of the deeper node
                update(1, 0, n-1, tin[deeper], tout[deeper], newW - oldW);
                
                // Update adjacency list
                for (auto &edge : adj[u]) {
                    if (edge.first == v) {
                        edge.second = newW;
                        break;
                    }
                }
                for (auto &edge : adj[v]) {
                    if (edge.first == u) {
                        edge.second = newW;
                        break;
                    }
                }
            } else {
                int x = query[1];
                long long dist = depth[x] + query(1, 0, n-1, tin[x]);
                result.push_back((int)dist);
            }
        }
        
        return result;
    }
};

class Solution:
    def treeQueries(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
        from collections import defaultdict
        
        adj = defaultdict(list)
        edge_weight = {}
        
        for u, v, w in edges:
            adj[u].append((v, w))
            adj[v].append((u, w))
            edge_weight[(min(u, v), max(u, v))] = w
        
        tin = [0] * (n + 1)
        tout = [0] * (n + 1)
        depth = [0] * (n + 1)
        timer = [0]
        
        def dfs(v, parent, d):
            tin[v] = timer[0]
            timer[0] += 1
            depth[v] = d
            for u, weight in adj[v]:
                if u != parent:
                    dfs(u, v, d + weight)
            tout[v] = timer[0] - 1
        
        dfs(1, -1, 0)
        
        # Segment tree with lazy propagation
        tree = [0] * (4 * n)
        lazy = [0] * (4 * n)
        
        def push(v, tl, tr):
            if lazy[v] != 0:
                tree[v] += lazy[v]
                if tl != tr:
                    lazy[v*2] += lazy[v]
                    lazy[v*2+1] += lazy[v]
                lazy[v] = 0
        
        def update(v, tl, tr, l, r, val):
            if l > r:
                return
            if l == tl and r == tr:
                lazy[v] += val
                push(v, tl, tr)
                return
            push(v, tl, tr)
            tm = (tl + tr) // 2
            update(v*2, tl, tm, l, min(r, tm), val)
            update(v*2+1, tm+1, tr, max(l, tm+1), r, val)
        
        def query(v, tl, tr, pos):
            if tl == tr:
                push(v, tl, tr)
                return tree[v]
            push(v, tl, tr)
            tm = (tl + tr) // 2
            if pos <= tm:
                return query(v*2, tl, tm, pos)
            else:
                return query(v*2+1, tm+1, tr, pos)
        
        result = []
        
        for q in queries:
            if q[0] == 1:
                u, v, new_w = q[1], q[2], q[3]
                old_w = edge_weight[(min(u, v), max(u, v))]
                edge_weight[(min(u, v), max(u, v))] = new_w
                
                deeper = u if depth[u] > depth[v] else v
                update(1, 0, n-1, tin[deeper], tout[deeper], new_w - old_w)
                
                # Update adjacency list
                for i, (neighbor, weight) in enumerate(adj[u]):
                    if neighbor == v:
                        adj[u][i] = (neighbor, new_w)
                        break
                for i, (neighbor, weight) in enumerate(adj[v]):
                    if neighbor == u:
                        adj[v][i] = (neighbor, new_w)
                        break
            else:
                x = q[1]
                dist = depth[x] + query(1, 0, n-1, tin[x])
                result.append(int(dist))
        
        return result

public class Solution {
    private List<List<(int, int)>> adj;
    private int[] tin, tout;
    private long[] depth;
    private long[] tree, lazy;
    private int timer;
    private int n;
    
    public int[] TreeQueries(int n, int[][] edges, int[][] queries) {
        this.n = n;
        adj = new List<List<(int, int)>>();
        for (int i = 0; i <= n; i++) {
            adj.Add(new List<(int, int)>());
        }
        
        tin = new int[n + 1];
        tout = new int[n + 1];
        depth = new long[n + 1];
        tree = new long[4 * n];
        lazy = new long[4 * n];
        timer = 0;
        
        var edgeWeight = new Dictionary<(int, int), int>();
        
        foreach (var edge in edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            adj[u].Add((v, w));
            adj[v].Add((u, w));
            edgeWeight[(Math.Min(u, v), Math.Max(u, v))] = w;
        }
        
        DFS(1, -1, 0);
        
        var result = new List<int>();
        
        foreach (var query in queries) {
            if (query[0] == 1) {
                int u = query[1], v = query[2], newW = query[3];
                int oldW = edgeWeight[(Math.Min(u, v), Math.Max(u, v))];
                edgeWeight[(Math.Min(u, v), Math.Max(u, v))] = newW;
                
                int deeper = depth[u] > depth[v] ? u : v;
                Update(1, 0, n - 1, tin[deeper], tout[deeper], newW - oldW);
                
                // Update adjacency list
                for (int i = 0; i < adj[u].Count; i++) {
                    if (adj[u][i].Item1 == v) {
                        adj[u][i] = (v, newW);
                        break;
                    }
                }
                for (int i = 0; i < adj[v].Count; i++) {
                    if (adj[v][i].Item1 == u) {
                        adj[v][i] = (u, newW);
                        break;
                    }
                }
            } else {
                int x = query[1];
                long dist = depth[x] + Query(1, 0, n - 1, tin[x]);
                result.Add((int)dist);
            }
        }
        
        return result.ToArray();
    }
    
    private void DFS(int v, int parent, long d) {
        tin[v] = timer++;
        depth[v] = d;
        foreach (var (u, weight) in adj[v]) {
            if (u != parent) {
                DFS(u, v, d + weight);
            }
        }
        tout[v] = timer - 1;
    }
    
    private void Push(int v, int tl, int tr) {
        if (lazy[v] != 0) {
            tree[v] += lazy[v];
            if (tl != tr) {
                lazy[v * 2] += lazy[v];
                lazy[v * 2 + 1] += lazy[v];
            }
            lazy[v] = 0;
        }
    }
    
    private void Update(int v, int tl, int tr, int l, int r, long val) {
        if (l > r) return;
        if (l == tl && r == tr) {
            lazy[v] += val;
            Push(v, tl, tr);
            return;
        }
        Push(v, tl, tr);
        int tm = (tl + tr) / 2;
        Update(v * 2, tl, tm, l, Math.Min(r, tm), val);
        Update(v * 2 + 1, tm + 1, tr, Math.Max(l, tm + 1), r, val);
    }
    
    private long Query(int v, int tl, int tr, int pos) {
        if (tl == tr) {
            Push(v, tl, tr);
            return tree[v];
        }
        Push(v, tl, tr);
        int tm = (tl + tr) / 2;
        if (pos <= tm)
            return Query(v * 2, tl, tm, pos);
        else
            return Query(v * 2 + 1, tm + 1, tr, pos);
    }
}

var treeQueries = function(n, edges, queries) {
    const adj = Array(n + 1).fill().map(() => []);
    const edgeMap = new Map();
    
    for (const [u, v, w] of edges) {
        adj[u].push([v, w]);
        adj[v].push([u, w]);
        const key = u < v ? `${u}-${v}` : `${v}-${u}`;
        edgeMap.set(key, w);
    }
    
    function dijkstra() {
        const dist = Array(n + 1).fill(Infinity);
        const pq = [[0, 1]];
        dist[1] = 0;
        
        while (pq.length > 0) {
            pq.sort((a, b) => a[0] - b[0]);
            const [d, u] = pq.shift();
            
            if (d > dist[u]) continue;
            
            for (const [v, weight] of adj[u]) {
                const key = u < v ? `${u}-${v}` : `${v}-${u}`;
                const actualWeight = edgeMap.get(key);
                
                if (dist[u] + actualWeight < dist[v]) {
                    dist[v] = dist[u] + actualWeight;
                    pq.push([dist[v], v]);
                }
            }
        }
        
        return dist;
    }
    
    const result = [];
    
    for (const query of queries) {
        if (query[0] === 1) {
            const [, u, v, w] = query;
            const key = u < v ? `${u}-${v}` : `${v}-${u}`;
            edgeMap.set(key, w);
        } else {
            const [, x] = query;
            const dist = dijkstra();
            result.push(dist[x]);
        }
    }
    
    return result;
};

在 LeetCode 上运行：力扣中文版 LeetCode 国际版

复杂度分析

复杂度	分析
时间复杂度	O((n + q) log n)
空间复杂度	O(n)

其中 n 是节点数，q 是查询数。DFS 预处理时间 O(n)，每个更新和查询操作都在线段树上执行，时间复杂度为 O(log n)。

题目描述#

解题思路#

代码实现#

复杂度分析#

题目描述

解题思路

代码实现

复杂度分析