Kimi K2.6 思考 LeetCode 3241. 标记所有节点需要的时间 Python3实现
LeetCode 3241. 标记所有节点需要的时间 — Python3 实现
题目概述
给定一棵无向树，节点编号 `0` 到 `n-1`。每个节点 `i` 被标记的规则：
- 奇数节点：相邻节点在时刻 `x-1` 被标记，则 `i` 在时刻 `x` 被标记（耗时 1）
- 偶数节点：相邻节点在时刻 `x-2` 被标记，则 `i` 在时刻 `x` 被标记（耗时 2）
求以每个节点为起点时，标记所有节点所需的最长时间。
核心思路
将问题转化为树形 DP + 换根（Re-rooting）：
1. 边权定义：从节点 `u` 到相邻节点 `v` 的"传播时间"取决于 `v` 的奇偶性
- `v` 为奇数：边权 = 1
- `v` 为偶数：边权 = 2
2. 两次 DFS：
- 第一次 DFS：计算每个节点子树内的最大深度和次大深度
- 第二次 DFS（换根）：计算每个节点往父节点方向的最长路径，得到最终答案
Python3 实现
```python
class Solution:
def timeTaken(self, edges: List[List[int]]) -> List[int]:
n = len(edges) + 1
g = [[] for _ in range(n)]
for x, y in edges:
g[x].append(y)
g[y].append(x)
# node[x] = (mx1, mx2, my)
# mx1: 子树 x 的最大深度, mx2: 次大深度, my: 最大深度对应的子节点
node = [[0, 0, 0] for _ in range(n)]
ans = [0] * n
def dfs(x: int, fa: int) -> int:
mx1 = mx2 = my = 0
for y in g[x]:
if y == fa:
continue
# 从 x 到 y 的边权: y 是偶数则 2, 奇数则 1
depth = dfs(y, x) + 2 - y % 2
if mx1 < depth:
mx2 = mx1
mx1 = depth
my = y
elif mx2 < depth:
mx2 = depth
node[x] = [mx1, mx2, my]
return mx1
def reroot(x: int, fa: int, from_: int):
ans[x] = max(from_, node[x][0])
for y in g[x]:
if y != fa:
# up1: 从 y 往上走到父节点 x 再继续往上
up1 = from_ + 2 - x % 2
# up2: 从 y 往上走到 x, 然后在 x 的其他子树中拐弯
up2 = (node[x][1] if y == node[x][2] else node[x][0]) + 2 - x % 2
reroot(y, x, max(up1, up2))
dfs(0, -1)
reroot(0, -1, 0)
return ans
```
关键逻辑说明
变量 含义
`node[x][0]` 以 `x` 为根的子树中，往子树方向的最大深度
`node[x][1]` 以 `x` 为根的子树中，往子树方向的次大深度
`node[x][2]` 最大深度对应的子节点编号
`from_` 从当前节点往父节点方向走的最长路径长度
`up1` 从子节点 `y` 往上经过 `x` 继续往上的路径
`up2` 从子节点 `y` 往上到 `x`，然后在 `x` 的其他子树中拐弯的路径
复杂度分析
- 时间复杂度：O(n)，每个节点访问常数次
- 空间复杂度：O(n)，递归栈和存储数组
示例验证
输入：`edges = [[0,1],[0,2]]`
- `i=0`：节点 1（奇数，耗时 1）→ `t=1`，节点 2（偶数，耗时 2）→ `t=2`，答案 = `2`
- `i=1`：节点 0（偶数，耗时 2）→ `t=2`，节点 2 需经过 0（偶数，耗时 2）→ `t=4`，答案 = `4`
- `i=2`：节点 0（偶数，耗时 2）→ `t=2`，节点 1（奇数，耗时 1）→ `t=3`，答案 = `3`
输出：`[2, 4, 3]` ✓