DeepSeek LeetCode 3585. 树中找到带权中位节点 Python3实现

解题思路核心思想利用树上倍增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 实现pythonfrom typing import Listclass 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] * ndist [0] * n# 迭代 DFS 预处理避免递归深度问题stack [(0, -1, 0, 0)] # (node, parent, depth, distance)while stack:u, p, d, w stack.pop()parent[0][u] pdepth[u] ddist[u] wfor v, wt in adj[u]:if v p:continuestack.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若不存在返回 -1i 0while k:if k 1:u parent[i][u]if u -1:breakk 1i 1return udef 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 ul 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 ufor k in range(LOG - 1, -1, -1):nxt parent[k][cur]if nxt ! -1 and dist[u] - dist[nxt] half:cur nxtreturn parent[0][cur]# 情况 2中位节点在 l - v 段else:need half - (dist[u] - dist[l])cur vfor k in range(LOG - 1, -1, -1):nxt parent[k][cur]if (nxt ! -1 and depth[nxt] depth[l] anddist[nxt] - dist[l] need):cur nxtreturn 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 无此限制无需额外处理。