这是 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 倍。

 

Logo

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

更多推荐