Kimi LeetCode 3544. 子树反转和 Python3实现

LeetCode 3544. 子树反转和 — Python3 实现思路概述这是一道 树形 DP 题目。核心观察点1. 反转的传递性若某个节点被反转其整个子树的所有节点值都会乘以 -1。因此可用一个状态 inv0/1表示当前子树是否处于“被反转”的符号下。2. 祖先链上的距离限制题目只限制存在祖先关系的两个反转节点之间的距离必须 ≥ k。非祖先关系的反转节点互不影响。因此只需在从根到叶的 DFS 路径上维护“距离上次反转还有多少步”。3. 状态设计dp(u, steps, inv) 表示以 u 为根的子树当前节点值为 inv ? -nums[u] : nums[u]且从 u 到最近祖先反转节点的距离为 steps若 steps k 则表示可以在 u 处再次反转时子树能得到的最大和。4. 状态转移- 不反转 u子节点继承当前符号 inv距离 1上限 k。- 反转 u仅当 steps k 时允许子节点符号翻转 1-inv距离重置为 1子节点到 u 的距离为 1。5. 复杂度状态数 O(n·k)每个状态遍历一次邻接表总时间 O(n·k)k ≤ 50n ≤ 5×10⁴完全可接受。6. 内存优化本题内存限制较紧Python 的 lru_cache 容易 MLE。采用一维数组扁平化存储 memo将 memo[u][inv][steps] 映射到 memo[u·2·(k1) inv·(k1) steps]可显著降低内存开销。---完整代码pythonimport sysfrom typing import List# 树深可能达到 5e4需要扩大递归限制sys.setrecursionlimit(200000)class Solution:def subtreeInversionSum(self, edges: List[List[int]], nums: List[int], k: int) - int:n len(edges) 1g [[] for _ in range(n)]for u, v in edges:g[u].append(v)g[v].append(u)parent [-1] * nINF 10 ** 15# 扁平化三维数组 memo[u][inv][steps]# 索引 u * 2 * (k1) inv * (k1) stepsmemo [INF] * (n * 2 * (k 1))def dfs(u: int, steps: int, inv: int) - int:idx u * 2 * (k 1) inv * (k 1) stepsif memo[idx] ! INF:return memo[idx]# 当前节点在 inv 状态下的实际值cur -nums[u] if inv else nums[u]no_flip cur# 不反转 u子节点继承符号距离 1到 k 封顶for v in g[u]:if v parent[u]:continueparent[v] uno_flip dfs(v, min(k, steps 1), inv)# 若距离上次反转已达 k可以选择反转 uif steps k:flip -curfor v in g[u]:if v parent[u]:continue# 反转 u 后子节点距离上次反转变为 1flip dfs(v, 1, 1 - inv)memo[idx] max(no_flip, flip)else:memo[idx] no_flipreturn memo[idx]# 从根出发stepsk 表示根节点随时可以反转return dfs(0, k, 0)---代码说明要点 说明steps 从当前节点向上到最近反转祖先的步数。steps k 意味着“可以再反转”。inv 0 表示当前子树未被反转1 表示被反转节点值取 -nums[u]。min(k, steps 1) 若不反转当前节点子节点到最近反转祖先的距离加 1但不超过 k超过 k 后统一视为 k表示“随时可以反转”。扁平化数组 避免 Python 嵌套列表的庞大对象开销将 O(n·k) 个状态压缩到一维 list实测可通过内存限制。递归限制 链状树深度可达 5×10⁴通过 sys.setrecursionlimit 防止栈溢出。---复杂度分析- 时间复杂度O(n · k)。每个状态 (u, steps, inv) 只计算一次每次遍历节点 u 的邻接点。- 空间复杂度O(n · k)。一维记忆化数组约 n × 2 × (k1) 个整数另加递归栈深度 O(n)。