以下是 LeetCode 3515. 带权树中的最短路径 的完整 Python3 实现。

核心思路

与 Java/Rust 版本完全一致:
1. DFS 欧拉序:将每个节点的子树映射到连续区间 `[tin, tout]`
2. 懒标记线段树:区间修改(子树整体加 `delta`)+ 单点查询(节点到根距离)
3. 边权更新:通过 `parent` 数组确定子节点,仅更新子树区间

完整 Python3 代码

```python
from typing import List
from collections import defaultdict

class LazySegmentTree:
    """Lazy Segment Tree supporting range add and point query."""
    
    def __init__(self, n: int):
        self.n = n
        self.tree = [0] * (4 * n)
        self.lazy = [0] * (4 * n)
    
    def _push(self, id: int, lo: int, hi: int) -> None:
        """Push lazy value down to children."""
        if self.lazy[id] == 0:
            return
        self.tree[id] += self.lazy[id]
        if lo != hi:
            self.lazy[id * 2] += self.lazy[id]
            self.lazy[id * 2 + 1] += self.lazy[id]
        self.lazy[id] = 0
    
    def add_range(self, l: int, r: int, val: int) -> None:
        """Add val to every element in range [l, r] (inclusive)."""
        self._add_range(1, 0, self.n - 1, l, r, val)
    
    def _add_range(self, id: int, lo: int, hi: int, l: int, r: int, val: int) -> None:
        self._push(id, lo, hi)
        if r < lo or l > hi:
            return
        if l <= lo and hi <= r:
            self.lazy[id] += val
            self._push(id, lo, hi)
            return
        mid = (lo + hi) // 2
        self._add_range(id * 2, lo, mid, l, r, val)
        self._add_range(id * 2 + 1, mid + 1, hi, l, r, val)
    
    def query(self, i: int) -> int:
        """Query value at index i."""
        return self._query(1, 0, self.n - 1, i)
    
    def _query(self, id: int, lo: int, hi: int, i: int) -> int:
        self._push(id, lo, hi)
        if lo == hi:
            return self.tree[id]
        mid = (lo + hi) // 2
        if i <= mid:
            return self._query(id * 2, lo, mid, i)
        return self._query(id * 2 + 1, mid + 1, hi, i)


class Solution:
    def treeQueries(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
        # Build adjacency list
        graph = [[] for _ in range(n + 1)]
        edge_weight = {}
        
        for u, v, w in edges:
            graph[u].append((v, w))
            graph[v].append((u, w))
            edge_weight[(min(u, v), max(u, v))] = w
        
        # Euler tour arrays
        tin = [0] * (n + 1)
        tout = [0] * (n + 1)
        parent = [0] * (n + 1)
        dist = [0] * (n + 1)
        
        # DFS to compute Euler tour, parent, and initial distances
        time = [0]
        
        def dfs(u: int, prev: int) -> None:
            tin[u] = time[0]
            time[0] += 1
            for v, w in graph[u]:
                if v == prev:
                    continue
                dist[v] = dist[u] + w
                parent[v] = u
                dfs(v, u)
            tout[u] = time[0] - 1
        
        dfs(1, 0)
        
        # Build segment tree with initial distances
        seg = LazySegmentTree(n)
        for i in range(1, n + 1):
            seg.add_range(tin[i], tin[i], dist[i])
        
        # Process queries
        ans = []
        for q in queries:
            if q[0] == 2:
                # Query: shortest path from root to node x
                x = q[1]
                ans.append(seg.query(tin[x]))
            else:
                # Update: change edge (u, v) weight to new_w
                u, v, new_w = q[1], q[2], q[3]
                key = (min(u, v), max(u, v))
                old_w = edge_weight[key]
                delta = new_w - old_w
                edge_weight[key] = new_w
                
                # Determine which node is the child (deeper in the tree)
                child = v if parent[v] == u else u
                
                # Update all nodes in child's subtree
                seg.add_range(tin[child], tout[child], delta)
        
        return ans
```

Python3 特有注意点

要点    说明    
嵌套函数 DFS    使用闭包捕获外部变量(`time`, `tin`, `tout` 等),无需手动传参    
线段树实现    用下划线前缀命名私有方法(`_push`, `_add_range`, `_query`)    
边权存储    `dict` 存储无向边,键为有序元组 `(min(u, v), max(u, v))`    
可变整数    `time` 用单元素列表 `[0]` 实现闭包内可变引用    

复杂度

- 时间复杂度:`O((n + q) log n)`
- 空间复杂度:`O(n)`

下载文件:[LeetCode3515_ShortestPathInWeightedTree.py](sandbox:///mnt/agents/output/LeetCode3515_ShortestPathInWeightedTree.py)

 

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐