当前位置：首页 > news >正文

Kimi LeetCode 3515. 带权树中的最短路径 Python3实现

news 2026/7/9 9:42:56

以下是 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)

查看全文

http://www.jsqmd.com/news/1153455/