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

Leetcode 3585. Find Weighted Median Node in Tree

news 来源：原创 2025/6/16 12:24:47

Leetcode 3585. Find Weighted Median Node in Tree
- 1. 解题思路
- 2. 代码实现

题目链接：3585. Find Weighted Median Node in Tree

1. 解题思路

这一题核心思路还是一个LCA算法，即最小公共父节点算法。

具体来说的话，对于每一个query给到的节点 $u, v$ ，我们都可以先通过LCA算法找到其最小公共父节点 $p$ ，然后我们可以提前在预处理过程中记录下每一个节点到根节点 $0$ 的距离为 $d_i$ ，此时不难发现三个节点 $u, v, p$ 到根节点的距离分别就是 $d_u, d_v, d_p$ ，而 $u, v$ 两节点之间的距离就是 $d_u+d_v-2d_p$ ，此时，我们不难找到 $uv$ 线段的中点。

此时，我们需要分类讨论一下：

如果其中点在线段右侧，即 $p v$ 这一段，那么我们要做的就是找到 $v$ 的父节点当中第一个距离大于等于 $d_v-(d_u+d_v-2d_p)/2$ 的节点；
如果其中点在线段左侧，即 $p u$ 这一段，那么我们要做的就是找到 $v$ 的父节点当中第一个距离小于 $d_v-(d_u+d_v-2d_p)/2$ 的节点；

而这个我们只需要稍微调整一下LCA的倍增过程即可快速实现。

2. 代码实现

给出python代码实现如下：

import math
from collections import deque
from typing import List, Tupleclass Tree:def __init__(self, n: int, edges: List[Tuple[int, int, int]], root: int = 0):self.n = nself.max_log = math.floor(math.log2(n)) + 1  # 最大跳跃步数的对数self.graph = [[] for _ in range(n)]self.distances = [0 for _ in range(n)]self.depth = [-1] * nself.parent = [[-1] * n for _ in range(self.max_log)]  # parent[k][i]: i 的第 2^k 级祖先# 构建邻接表for u, v, w in edges:self.graph[u].append((v, w))self.graph[v].append((u, w))# 预处理深度和祖先表self._bfs(root)def _bfs(self, root: int):"""BFS 初始化深度和直接父节点（即 2^0 级祖先）"""queue = deque([root])self.depth[root] = 0self.distances[root] = 0self.parent[0][root] = -1  # 根节点无父节点while queue:u = queue.popleft()for v, w in self.graph[u]:if v == self.parent[0][u]:continueself.depth[v] = self.depth[u] + 1self.distances[v] = self.distances[u] + wself.parent[0][v] = uqueue.append(v)# 递推计算 2^k 级祖先for k in range(1, self.max_log):for i in range(self.n):if self.parent[k-1][i] != -1:self.parent[k][i] = self.parent[k-1][self.parent[k-1][i]]def query(self, u: int, v: int) -> int:"""查询节点 u 和 v 的最近公共祖先"""# 确保 u 是深度较大的节点if self.depth[u] < self.depth[v]:u, v = v, u# 将 u 上跳到与 v 同深度diff = self.depth[u] - self.depth[v]k = 0while diff:if diff & 1:u = self.parent[k][u]diff >>= 1k += 1if u == v:return u# 同步上跳，寻找最近公共祖先for k in range(self.max_log - 1, -1, -1):if self.parent[k][u] != self.parent[k][v]:u = self.parent[k][u]v = self.parent[k][v]return self.parent[0][u]def query_distance(self, u):return self.distances[u]def query_closest_parent(self, u: int, d: float):h = self.max_log-1while h >= 0:if self.parent[h][u] != -1 and self.distances[self.parent[h][u]] >= d:u = self.parent[h][u]h -= 1return uclass Solution:def findMedian(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]:tree = Tree(n, edges, 0)def query(u, v):p = tree.query(u, v)du, dv, dp = tree.query_distance(u), tree.query_distance(v), tree.query_distance(p)d1, d2 = du-dp, dv-dpd = (d1+d2) / 2if d <= d2:return tree.query_closest_parent(v, dv-d)else:w = tree.query_closest_parent(u, du-d)dw = tree.query_distance(w)return tree.parent[0][w] if du-dw != d else wreturn [query(u, v) for u, v in queries]