Kimi LeetCode 3553. 包含要求路径的最小带权子图 II Python3实现
这是 LeetCode 3553 的 Python3 实现,核心思路与 Rust 版一致:
> 答案 = (dist(src1, src2) + dist(src1, dest) + dist(src2, dest)) // 2
即三个点的最小连通子树边权和,等于它们两两树上距离之和的一半。求树上距离需要 LCA 倍增。
```python
from typing import List
from collections import deque
class Solution:
def minimumWeight(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
n = len(edges) + 1
graph = [[] for _ in range(n)]
for u, v, w in edges:
graph[u].append((v, w))
graph[v].append((u, w))
LOG = n.bit_length()
up = [[0] * LOG for _ in range(n)] # 倍增祖先表
depth = [0] * n
dist = [0] * n # 到根节点 0 的距离
# BFS 预处理:深度、距离、直接父节点
q = deque([0])
visited = [False] * n
visited[0] = True
while q:
u = q.popleft()
for v, w in graph[u]:
if not visited[v]:
visited[v] = True
up[v][0] = u
depth[v] = depth[u] + 1
dist[v] = dist[u] + w
q.append(v)
# 构建倍增表
for j in range(1, LOG):
for i in range(n):
up[i][j] = up[up[i][j - 1]][j - 1]
# LCA 查询
def lca(u: int, v: int) -> int:
if depth[u] > depth[v]:
u, v = v, u
# 将 v 提升到 u 的深度
diff = depth[v] - depth[u]
for i in range(LOG):
if diff >> i & 1:
v = up[v][i]
if u == v:
return u
# 从大到小同时上跳
for i in range(LOG - 1, -1, -1):
if up[u][i] != up[v][i]:
u = up[u][i]
v = up[v][i]
return up[u][0]
# 树上距离
def distance(u: int, v: int) -> int:
w = lca(u, v)
return dist[u] + dist[v] - 2 * dist[w]
# 回答查询
ans = []
for src1, src2, dest in queries:
d12 = distance(src1, src2)
d1d = distance(src1, dest)
d2d = distance(src2, dest)
ans.append((d12 + d1d + d2d) // 2)
return ans
```
思路说明
步骤 说明
建图 无向带权树,邻接表存储
BFS 预处理 以 `0` 为根,求每个节点的深度 `depth[]`、到根距离 `dist[]`、直接父节点 `up[v][0]`
倍增表 `up[v][j]` 表示 `v` 的 2^j 级祖先,O(n \log n) 预处理
LCA 每次查询 O(\log n):先统一深度,再同时向上跳
树上距离 \text{dist}(u, v) = \text{dist}[u] + \text{dist}[v] - 2 \cdot \text{dist}[\text{lca}(u,v)]
最终答案 对每个查询,(\text{dist}(src1, src2) + \text{dist}(src1, dest) + \text{dist}(src2, dest)) // 2
复杂度
- 时间:O(n \log n + q \log n),其中 q 为查询数
- 空间:O(n \log n)
> 为什么公式成立? 树中三个点的最小连通子树,恰好是三条两两路径的并集。该子树内的每条边在三条路径中恰好被经过 2 次(因为树无环,每条边将树分成两部分,三个点必有两点在同侧),所以三条距离之和等于子树边权和的 2 倍。
更多推荐



所有评论(0)