Kimi LeetCode 3367. 移除边之后的权重最大和 Python3实现
LeetCode 3367. 移除边之后的权重最大和 — Python3 实现题目描述给定一棵无向树n 个节点n-1 条边删除部分边使得每个节点至多与 k 个其他节点相连求剩余边权重的最大和。核心思路树形 DP 贪心以节点 0 为根进行 DFS。对于每个节点 u定义两个状态状态 含义dp_k 节点 u 最多连接 k 条边 到子节点的最大权重和dp_k1 节点 u 最多连接 k-1 条边 到子节点的最大权重和预留一条给父节点对于子节点 v边权为 w- 不选边 (u,v)子树贡献 dp_k(v)- 选边 (u,v)子树贡献 w dp_k1(v)选择这条边的收益增量为delta w dp_k1(v) - dp_k(v)对每个节点收集所有正收益增量按降序排序贪心取前 k 个或 k-1 个。Python3 实现pythonfrom typing import List, Tupleclass Solution:def maximizeSumOfWeights(self, edges: List[List[int]], k: int) - int:n len(edges) 1# 构建邻接表g [[] for _ in range(n)]for u, v, w in edges:g[u].append((v, w))g[v].append((u, w))def dfs(u: int, fa: int) - Tuple[int, int]:返回:- 第一个值: 节点 u 最多选 k 条边连子节点的最大权重和- 第二个值: 节点 u 最多选 k-1 条边连子节点的最大权重和预留一条给父节点s 0 # 所有子节点都不选边时的基础和t [] # 正收益增量列表for v, w in g[u]:if v fa:continuea, b dfs(v, u) # a dp_k(v), b dp_k1(v)s a # 默认不选边 (u,v)# 选择边 (u,v) 的收益增量d w b - aif d 0:t.append(d)# 按收益降序排序贪心选最优的边t.sort(reverseTrue)# 选 k-1 条边预留一条给父节点sum_k1 s sum(t[:k - 1])# 选 k 条边sum_k s sum(t[:k])return sum_k, sum_k1x, y dfs(0, -1)# 根节点没有父节点两种状态都合法取最大值return max(x, y)另一种写法使用 heapq.nlargestpythonimport heapqfrom typing import List, Tupleclass Solution:def maximizeSumOfWeights(self, edges: List[List[int]], k: int) - int:n len(edges) 1g [[] for _ in range(n)]for u, v, w in edges:g[u].append((v, w))g[v].append((u, w))def dfs(u: int, prev: int) - Tuple[int, int]:weight_sum 0diffs []for v, w in g[u]:if v prev:continuesub_k1, sub_k dfs(v, u)weight_sum sub_k# 选择边 (u,v) 的收益增量diffs.append(max(0, sub_k1 - sub_k w))# 取前 k 个 / 前 k-1 个最大收益return (weight_sum sum(heapq.nlargest(k, diffs)),weight_sum sum(heapq.nlargest(k - 1, diffs)))return dfs(0, -1)[1]复杂度分析指标 复杂度时间 O(n log n) — 每个节点的子节点收益排序或用 heapq空间 O(n) — 邻接表 递归栈关键点说明1. 收益增量 delta w dp_k1(v) - dp_k(v)表示选择边 (u,v) 相比不选能额外获得的收益。如果为负则不选更优。2. 两种状态设计- dp_k节点 u 可以向子节点连最多 k 条边- dp_k1节点 u 只能向子节点连最多 k-1 条边预留一个位置给父节点3. 根节点处理根节点没有父节点所以最终答案取 max(x, y) 均可。4. Python 递归深度n 最大为 10^5Python 默认递归深度限制为 1000需要在提交时设置 sys.setrecursionlimit(200000)。完整可运行版本含递归深度设置pythonimport sysfrom typing import List, Tuple# 提升递归深度限制sys.setrecursionlimit(200000)class Solution:def maximizeSumOfWeights(self, edges: List[List[int]], k: int) - int:n len(edges) 1g [[] for _ in range(n)]for u, v, w in edges:g[u].append((v, w))g[v].append((u, w))def dfs(u: int, fa: int) - Tuple[int, int]:s 0t []for v, w in g[u]:if v fa:continuea, b dfs(v, u)s ad w b - aif d 0:t.append(d)t.sort(reverseTrue)return s sum(t[:k]), s sum(t[:k - 1])x, y dfs(0, -1)return max(x, y)验证示例示例 1edges [[0,1,4],[0,2,2],[2,3,12],[2,4,6]], k 2 → 22- 节点 2 有 3 条边连 0,3,4需删除一条- 删除边 (0,2,2)收益最小保留 (2,3,12) 和 (2,4,6)- 加上 (0,1,4)总和 4 12 6 22示例 2edges [[0,1,5],[1,2,10],[0,3,15],[3,4,20],[3,5,5],[0,6,10]], k 3 → 65- 没有节点超过 3 条边全部保留- 总和 5 10 15 20 5 10 65