DeepSeek LeetCode 3559. 给边赋权值的方案数 II Python3实现

稚辉君.MCA.P9.JAVA

158人浏览 · 2026-07-14 08:41:19

稚辉君.MCA.P9.JAVA · 2026-07-14 08:41:19 发布

根据题目要求，对于树中任意两个节点 u 和 v，路径长度为 d（边数），使路径总代价为奇数的赋值方案数为 2^(d-1)。因此核心是快速求树上两点距离。

Python3 实现

```python
class Solution:
MOD = 10**9 + 7

def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
n = len(edges) + 1

# 1. 建图（无向树）
graph = [[] for _ in range(n)]
for u, v in edges:
u -= 1
v -= 1
graph[u].append(v)
graph[v].append(u)

# 2. 预处理深度和倍增表
LOG = (n).bit_length() # ceil(log2(n)) + 1
depth = [0] * n
up = [[0] * LOG for _ in range(n)]

# DFS 构建倍增表
stack = [(0, -1)] # (node, parent)
order = [] # 后序遍历顺序
parent = [-1] * n

while stack:
node, par = stack.pop()
parent[node] = par
order.append(node)
for nei in graph[node]:
if nei != par:
depth[nei] = depth[node] + 1
stack.append((nei, node))

# 初始化 up[0]
for i in range(n):
up[i][0] = parent[i] if parent[i] != -1 else i

# 构建倍增表
for j in range(1, LOG):
for i in range(n):
up[i][j] = up[up[i][j-1]][j-1]

# 3. 预处理 2 的幂次
pow2 = [1] * (n + 1)
for i in range(1, n + 1):
pow2[i] = (pow2[i-1] * 2) % self.MOD

# 4. 处理每个查询
ans = []
for u, v in queries:
u -= 1
v -= 1
lca = self._find_lca(u, v, depth, up, LOG)
dist = depth[u] + depth[v] - 2 * depth[lca]
ans.append(0 if dist == 0 else pow2[dist - 1])

return ans

def _find_lca(self, u, v, depth, up, LOG):
"""求 LCA"""
# 确保 u 是较浅的节点
if depth[u] > depth[v]:
u, v = v, u

# 将 v 提升到与 u 同深度
diff = depth[v] - depth[u]
for j in range(LOG):
if (diff >> j) & 1:
v = up[v][j]

if u == v:
return u

# 同时提升
for j in range(LOG - 1, -1, -1):
if up[u][j] != up[v][j]:
u = up[u][j]
v = up[v][j]

return up[u][0]
```

优化版本（使用 BFS 替代 DFS）

```python
from collections import deque

class Solution:
MOD = 10**9 + 7

def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
n = len(edges) + 1

# 建图
graph = [[] for _ in range(n)]
for u, v in edges:
u, v = u - 1, v - 1
graph[u].append(v)
graph[v].append(u)

# BFS 预处理深度和父节点
LOG = (n).bit_length()
depth = [0] * n
up = [[0] * LOG for _ in range(n)]

q = deque([0])
visited = [False] * n
visited[0] = True
parent = [-1] * n

while q:
node = q.popleft()
for nei in graph[node]:
if not visited[nei]:
visited[nei] = True
depth[nei] = depth[node] + 1
parent[nei] = node
q.append(nei)

# 初始化倍增表
for i in range(n):
up[i][0] = parent[i] if parent[i] != -1 else i

for j in range(1, LOG):
for i in range(n):
up[i][j] = up[up[i][j-1]][j-1]

# 预处理 2 的幂
pow2 = [1] * (n + 1)
for i in range(1, n + 1):
pow2[i] = (pow2[i-1] * 2) % self.MOD

# 处理查询
def lca(u, v):
if depth[u] > depth[v]:
u, v = v, u

# 提升 v
diff = depth[v] - depth[u]
j = 0
while diff:
if diff & 1:
v = up[v][j]
diff >>= 1
j += 1

if u == v:
return u

for j in range(LOG - 1, -1, -1):
if up[u][j] != up[v][j]:
u = up[u][j]
v = up[v][j]

return up[u][0]

ans = []
for u, v in queries:
u, v = u - 1, v - 1
ancestor = lca(u, v)
dist = depth[u] + depth[v] - 2 * depth[ancestor]
ans.append(0 if dist == 0 else pow2[dist - 1])

return ans
```

更简洁的实现（使用递归 DFS）

```python
class Solution:
MOD = 10**9 + 7

def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
n = len(edges) + 1

# 建图
g = [[] for _ in range(n)]
for u, v in edges:
u, v = u - 1, v - 1
g[u].append(v)
g[v].append(u)

# DFS 预处理
LOG = n.bit_length()
depth = [0] * n
up = [[0] * LOG for _ in range(n)]

def dfs(node, parent):
up[node][0] = parent if parent != -1 else node
for j in range(1, LOG):
up[node][j] = up[up[node][j-1]][j-1]
for nei in g[node]:
if nei != parent:
depth[nei] = depth[node] + 1
dfs(nei, node)

dfs(0, -1)

# LCA 函数
def get_lca(u, v):
if depth[u] > depth[v]:
u, v = v, u
# 提升 v
diff = depth[v] - depth[u]
j = 0
while diff:
if diff & 1:
v = up[v][j]
diff >>= 1
j += 1
if u == v:
return u
for j in range(LOG - 1, -1, -1):
if up[u][j] != up[v][j]:
u = up[u][j]
v = up[v][j]
return up[u][0]

# 预处理 2 的幂
pow2 = [1] * (n + 1)
for i in range(1, n + 1):
pow2[i] = (pow2[i-1] * 2) % self.MOD

# 处理查询
return [
0 if dist == 0 else pow2[dist - 1]
for u, v in queries
for dist in [depth[u-1] + depth[v-1] - 2 * depth[get_lca(u-1, v-1)]]
]
```