• leetcode 3559. Number of Ways to Assign Edge Weights II
  • 1. 解題思路
  • 2. 代碼實現

• 題目鏈接：3559. Number of Ways to Assign Edge Weights II

1. 解題思路

這一題是題目3558. Number of Ways to Assign Edge Weights I的進階版本。

對于題目3558來說，其核心就是找出樹的最大深度，然后就是一個數學題了，要使得這一條路徑上的邊的總和為奇數，那么其可選的方式為：
$$N=\sum _{i=1}^{i\le n,i=1,3,5,?}{C}_n^i=2^{n?1}$$

因此，我們可以直接計算出其結果。

而對于本題，問題增加的難度在于需要快速地任意給到的兩個點$u,v$，找出這兩點間的連接線段的長度，而這個的話事實上和題目3553. Minimum Weighted Subgraph With the Required Paths II非常相近，事實上我們只需要分別求出這兩個點$u,v$的深度以及他們的最小公共父節點$p$的深度，那么$u,v$間的距離$d$就可以快速通過下述公式計算得到：
$$d=h_u+h_v?2\times h_p$$

剩下的就是如何求出$u,v$的最小公共父節點了，而這個事實上就是一個經典算法了，這里就不具體展開了，后續有時間的話可能會系統整理一下這個算法然后發一篇博客作為備忘好了。

2. 代碼實現

給出python代碼實現如下：

MOD = 10**9+7class LCA:def __init__(self, root, tree):self.max_level = 20  # 根據樹的高度調整self.parent = defaultdict(lambda: [-1]*self.max_level)self.depth = {}self.preprocess(root, tree)def preprocess(self, root, tree):stack = [(root, -1, 0)]  # (node, parent, depth)while stack:node, par, d = stack.pop()self.depth[node] = dself.parent[node][0] = parfor k in range(1, self.max_level):if self.parent[node][k-1] != -1:self.parent[node][k] = self.parent[self.parent[node][k-1]][k-1]for child in tree[node]:if child != par:stack.append((child, node, d+1))def query(self, u, v):if self.depth[u] < self.depth[v]:u, v = v, u# 對齊深度for k in range(self.max_level-1, -1, -1):if self.depth[u] - (1 << k) >= self.depth[v]:u = self.parent[u][k]if u == v:return u# 同步跳轉for k in range(self.max_level-1, -1, -1):if self.parent[u][k] != -1 and self.parent[u][k] != self.parent[v][k]:u = self.parent[u][k]v = self.parent[v][k]return self.parent[u][0]class Solution:def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:graph = defaultdict(list)for u, v in edges:graph[u].append(v)graph[v].append(u)depth = defaultdict(int)tree = defaultdict(list)def dfs(u, p):nonlocal depth, treeif u == 1:depth[u] = 0else:depth[u] = depth[p] + 1for v in graph[u]:if v == p:continuetree[u].append(v)dfs(v, u)returndfs(1, 0)lca = LCA(1, tree)def query(u, v):p = lca.query(u, v)d = depth[u] + depth[v] - 2*depth[p]return pow(2, d-1, MOD) if d > 0 else 0return [query(u, v) for u, v in queries]

提交代碼評測得到：耗時2773ms，占用內存137.3MB。