解题思路
核心思想:利用树上倍增(LCA)+ 二分跳跃在 O(log n) 时间内回答每个查询。
1. 预处理:
· 以节点 0 为根,DFS 计算每个节点的深度 depth[]、到根的距离 dist[](边权和),以及倍增祖先表 parent[k][u](表示 u 的 2^k 级祖先)。
· 树中边权为正,因此路径距离单调递增。
2. 查询处理((u, v)):
· 若 u == v,答案就是 u。
· 求 u 和 v 的 LCA l,计算路径总权值 total = dist[u] + dist[v] - 2*dist[l],取半阈值 half = (total + 1) // 2(向上取整)。
· 情况一:中位节点在 u → l 这一段(即 dist[u] - dist[l] >= half)。从 u 向上跳,找到最后一个满足 dist[u] - dist[ancestor] < half 的祖先,答案就是该祖先的父节点。
· 情况二:中位节点在 l → v 这一段。先计算 need = half - (dist[u] - dist[l]),表示从 l 向下还需走的权值。从 v 向上跳,找到最浅(最靠近 l)的满足 dist[ancestor] - dist[l] >= need 的祖先,该祖先即为答案。
---
Python3 实现
```python
from typing import List
class Solution:
def findMedian(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
# 建图(邻接表)
adj = [[] for _ in range(n)]
for u, v, w in edges:
adj[u].append((v, w))
adj[v].append((u, w))
LOG = (n).bit_length() + 1 # 倍增表大小
parent = [[-1] * n for _ in range(LOG)]
depth = [0] * n
dist = [0] * n
# 迭代 DFS 预处理(避免递归深度问题)
stack = [(0, -1, 0, 0)] # (node, parent, depth, distance)
while stack:
u, p, d, w = stack.pop()
parent[0][u] = p
depth[u] = d
dist[u] = w
for v, wt in adj[u]:
if v == p:
continue
stack.append((v, u, d + 1, w + wt))
# 构建倍增表
for k in range(1, LOG):
prev = parent[k - 1]
curr = parent[k]
for i in range(n):
anc = prev[i]
curr[i] = prev[anc] if anc != -1 else -1
# ---------- 辅助函数 ----------
def kth_ancestor(u: int, k: int) -> int:
"""返回 u 的第 k 级祖先(k >= 0),若不存在返回 -1"""
i = 0
while k:
if k & 1:
u = parent[i][u]
if u == -1:
break
k >>= 1
i += 1
return u
def lca(u: int, v: int) -> int:
"""返回 u 和 v 的最近公共祖先"""
if depth[u] < depth[v]:
u, v = v, u
# 提升 u 至与 v 同深度
u = kth_ancestor(u, depth[u] - depth[v])
if u == v:
return u
# 一起向上跳
for k in range(LOG - 1, -1, -1):
if parent[k][u] != parent[k][v]:
u = parent[k][u]
v = parent[k][v]
return parent[0][u]
def weighted_median(u: int, v: int) -> int:
if u == v:
return u
l = lca(u, v)
total = dist[u] + dist[v] - 2 * dist[l]
half = (total + 1) // 2 # 向上取整
# 情况 1:中位节点在 u -> l 段
if dist[u] - dist[l] >= half:
cur = u
for k in range(LOG - 1, -1, -1):
nxt = parent[k][cur]
if nxt != -1 and dist[u] - dist[nxt] < half:
cur = nxt
return parent[0][cur]
# 情况 2:中位节点在 l -> v 段
else:
need = half - (dist[u] - dist[l])
cur = v
for k in range(LOG - 1, -1, -1):
nxt = parent[k][cur]
if (nxt != -1 and depth[nxt] >= depth[l] and
dist[nxt] - dist[l] >= need):
cur = nxt
return cur
# 处理所有查询
return [weighted_median(u, v) for u, v in queries]
```
---
复杂度分析
· 时间复杂度:预处理 O(n log n),每个查询 O(log n),总复杂度 O((n + q) log n)。
· 空间复杂度:邻接表 O(n),倍增表 O(n log n),总复杂度 O(n log n)。
注意:边权和路径总权值可能超过 32 位整数范围,Python 的 int 无此限制,无需额外处理。